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.

4 files changed, 71 insertions, 8 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index a9d53162c..3a7df8e7a 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -1039,6 +1039,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);
 
diff --git a/source/slang/slang-emit-spirv.cpp b/source/slang/slang-emit-spirv.cpp
index 90d0db5d7..b0da7e4a6 100644
--- a/source/slang/slang-emit-spirv.cpp
+++ b/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,
diff --git a/source/slang/slang-ir-inst-defs.h b/source/slang/slang-ir-inst-defs.h
index 33baefa51..6d4f7a2ca 100644
--- a/source/slang/slang-ir-inst-defs.h
+++ b/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 */
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index fb06863d4..4c7272755 100644
--- a/source/slang/slang-ir-insts.h
+++ b/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;