Add intrinsics for aligned load/store. (#5736)

* Add intrinsics for aligned load/store.

* Fix.

* Update comment.

* Implement aligned load/store as intrinsic_op.

* Fix.

* Add proposal doc.

* fix typo.

Author: csyonghe

docs/proposals/013-aligned-load-store.md

@@ -0,0 +1,58 @@
+SP #013: Aligned load store
+=========================================
+
+Status: Experimental
+
+Implementation: [PR 5736](https://github.com/shader-slang/slang/pull/5736)
+
+Author: Yong He (yhe@nvidia.com)
+
+Reviewer: 
+
+Introduction
+----------
+
+On many architectures, aligned vector loads (e.g. loading a float4 with 16 byte alignment) is often more efficient than ordinary unaligned loads. Slang's pointer type does not encode any additional alignment info, and all pointer read/writes are by default assuming the alignment of the underlying pointee type, which is 4 bytes for float4 vectors. This means that loading from a `float4*` will result in unaligned load instructions.
+
+This proposal attempts to provide a way for performance sensitive code to specify an aligned load/store through Slang pointers.
+
+
+Proposed Approach
+------------
+
+We propose to add intrinsic functions to perform aligned load/store through a pointer:
+
+```
+T loadAligned<int alignment, T>(T* ptr);
+void storeAligned<int alignment, T>(T* ptr, T value);
+```
+
+Example:
+
+```
+uniform float4* data;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    var v = loadAligned<8>(data);
+    storeAligned<16>(data+1, v);
+}
+```
+
+Related Work
+------------
+
+### GLSL ###
+
+GLSL supports the `align` layout on a `buffer_reference` block to specify the alignment of the buffer pointer.
+
+### SPIRV ###
+
+In SPIRV, the alignment can either be encoded as a decoration on the pointer type, or as a memory operand on the OpLoad and OpStore operations.
+
+### Other Languages ###
+
+Most C-like languages allow users to put additional attributes on types to specify the alignment of the type. All loads/stores through pointers of the type will use the alignment.
+
+Instead of introducing type modifiers on data or pointer types, Slang should explicitly provide a `loadAligned` and `storeAligned` intrinsic functions to leads to `OpLoad` and `OpStore` with the `Aligned` memory operand when generating SPIRV. This way we don't have to deal with the complexity around rules of handling type coercion between modified/unmodified types and recalculate alignment for pointers representing an access chain. Developers writing performance sentisitive code can always be assured that the alignment specified on each critical load or store will be assumed, without having to work backwards through type modifications and thinking about the typing rules associated with such modifiers.

source/slang/core.meta.slang

@@ -1038,6 +1038,47 @@ struct Ptr
     }
 };
 
+//@hidden:
+__intrinsic_op($(kIROp_AlignedAttr))
+void __align_attr(int alignment);
+
+__intrinsic_op($(kIROp_Load))
+T __load_aligned<T, U>(T* ptr, U alignmentAttr);
+
+__intrinsic_op($(kIROp_Store))
+void __store_aligned<T, U>(T* ptr, T value, U alignmentAttr);
+
+//@public:
+
+/// Load a value from a pointer with a known alignment.
+/// Aligned loads are more efficient than unaligned loads on some platforms.
+/// @param alignment The alignment of the load operation.
+/// @param ptr The pointer to load from.
+/// @return The value loaded from the pointer.
+/// @remarks When targeting SPIRV, this function maps to an `OpLoad` instruction with the `Aligned` memory operand.
+/// The functions maps to normal load operation on other targets.
+///
+[__NoSideEffect]
+[ForceInline]
+T loadAligned<int alignment, T>(T* ptr)
+{
+    return __load_aligned(ptr, __align_attr(alignment));
+}
+
+/// Store a value to a pointer with a known alignment.
+/// Aligned stores are more efficient than unaligned stores on some platforms.
+/// @param alignment The alignment of the store operation.
+/// @param ptr The pointer to store value to.
+/// @param value The value to store.
+/// @remarks When targeting SPIRV, this function maps to an `OpStore` instruction with the `Aligned` memory operand.
+/// The functions maps to normal store operation on other targets.
+///
+[ForceInline]
+void storeAligned<int alignment, T>(T* ptr, T value)
+{
+    __store_aligned(ptr, value, __align_attr(alignment));
+}
+
 //@hidden:
 __intrinsic_op($(kIROp_Load))
 T __load<T, let addrSpace : uint64_t>(Ptr<T, addrSpace> ptr);

source/slang/slang-emit-spirv.cpp

@@ -5975,10 +5975,17 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                            SpvStorageClassPhysicalStorageBuffer)
         {
             IRSizeAndAlignment sizeAndAlignment;
-            getNaturalSizeAndAlignment(
-                m_targetProgram->getOptionSet(),
-                ptrType->getValueType(),
-                &sizeAndAlignment);
+            if (auto alignedAttr = inst->findAttr<IRAlignedAttr>())
+            {
+                sizeAndAlignment.alignment = (int)getIntVal(alignedAttr->getAlignment());
+            }
+            else
+            {
+                getNaturalSizeAndAlignment(
+                    m_targetProgram->getOptionSet(),
+                    ptrType->getValueType(),
+                    &sizeAndAlignment);
+            }
             return emitOpLoadAligned(
                 parent,
                 inst,
@@ -5999,10 +6006,17 @@ struct SPIRVEmitContext : public SourceEmitterBase, public SPIRVEmitSharedContex
                            SpvStorageClassPhysicalStorageBuffer)
         {
             IRSizeAndAlignment sizeAndAlignment;
-            getNaturalSizeAndAlignment(
-                m_targetProgram->getOptionSet(),
-                ptrType->getValueType(),
-                &sizeAndAlignment);
+            if (auto alignedAttr = inst->findAttr<IRAlignedAttr>())
+            {
+                sizeAndAlignment.alignment = (int)getIntVal(alignedAttr->getAlignment());
+            }
+            else
+            {
+                getNaturalSizeAndAlignment(
+                    m_targetProgram->getOptionSet(),
+                    ptrType->getValueType(),
+                    &sizeAndAlignment);
+            }
             return emitOpStoreAligned(
                 parent,
                 inst,

source/slang/slang-ir-inst-defs.h

@@ -1250,6 +1250,7 @@ INST_RANGE(Layout, VarLayout, EntryPointLayout)
     INST(SNormAttr, snorm, 0, HOISTABLE)
     INST(NoDiffAttr, no_diff, 0, HOISTABLE)
     INST(NonUniformAttr, nonuniform, 0, HOISTABLE)
+    INST(AlignedAttr, Aligned, 1, HOISTABLE)
 
     /* SemanticAttr */
         INST(UserSemanticAttr, userSemantic, 2, HOISTABLE)
@@ -1260,6 +1261,7 @@ INST_RANGE(Layout, VarLayout, EntryPointLayout)
         INST(VarOffsetAttr, offset, 2, HOISTABLE)
     INST_RANGE(LayoutResourceInfoAttr, TypeSizeAttr, VarOffsetAttr)
     INST(FuncThrowTypeAttr, FuncThrowType, 1, HOISTABLE)
+    
 INST_RANGE(Attr, PendingLayoutAttr, FuncThrowTypeAttr)
 
 /* Liveness */

source/slang/slang-ir-insts.h

@@ -2389,6 +2389,12 @@ struct IRCall : IRInst
     void setArg(UInt index, IRInst* arg) { setOperand(index + 1, arg); }
 };
 
+struct IRAlignedAttr : IRAttr
+{
+    IR_LEAF_ISA(AlignedAttr)
+    IRInst* getAlignment() { return getOperand(0); }
+};
+
 struct IRLoad : IRInst
 {
     IRUse ptr;

tests/spirv/aligned-load-store.slang

@@ -0,0 +1,13 @@
+//TEST:SIMPLE(filecheck=CHECK): -target spirv
+
+// CHECK: OpLoad {{.*}} Aligned 8
+// CHECK: OpStore {{.*}} Aligned 16
+
+uniform float4* data;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    var v = loadAligned<8>((float2x4*)data);
+    storeAligned<16>((float2x4*)data+1, v);
+}