From 00d1206c59da83f92266dfbb98524ae3e3de86db Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Thu, 28 May 2026 14:10:33 +0000
Subject: [PATCH] Fix dropped alignment operand in byteAddressBufferLoad
 specialization

specializeFuncsForBufferLoadArgs lumped kIROp_ByteAddressBufferLoad in
with the generic two-operand "element access" instructions (GetElement,
StructuredBufferLoad, etc.) in both getCallInfoForArg and
getSpecializedValueForArg. byteAddressBufferLoad actually has three
operands -- (buffer, offset, alignment) -- so the rewritten load inside
the specialized callee was constructed with only two operands.

Downstream, legalizeByteAddressBufferOps::processLoad unconditionally
reads operand index 2 to recover the alignment, hitting

    assert failure: slang-ir.h(710): index < getOperandCount()

during SPIR-V emission. A minimal reproducer is a struct large enough
to trip the buffer-load specialization threshold (e.g. two float4x4s
plus a padded float3) loaded via Buf.Load<T>() and passed by value into
a helper function.

Give byteAddressBufferLoad dedicated handling in both functions so the
alignment operand survives specialization, and add a regression test
under tests/optimization that exercises the failing pattern on SPIR-V.
---
 .../slang-ir-specialize-function-call.cpp     | 46 ++++++++++++++++++-
 .../buffer-load-specialize-byte-address.slang | 37 +++++++++++++++
 2 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 tests/optimization/buffer-load-specialize-byte-address.slang

diff --git a/source/slang/slang-ir-specialize-function-call.cpp b/source/slang/slang-ir-specialize-function-call.cpp
index 9e7afb46aa3..96b36cb5367 100644
--- a/source/slang/slang-ir-specialize-function-call.cpp
+++ b/source/slang/slang-ir-specialize-function-call.cpp
@@ -617,6 +617,28 @@ struct FunctionParameterSpecializationContext
 
             ioInfo.newArgs.add(oldIndex);
         }
+        else if (oldArg->getOp() == kIROp_ByteAddressBufferLoad)
+        {
+            // Like the element-access case above, but with an extra `alignment` operand.
+            auto oldBuffer = oldArg->getOperand(0);
+            auto oldOffset = oldArg->getOperand(1);
+            auto oldAlignment = oldArg->getOperand(2);
+
+            getCallInfoForArg(ioInfo, oldBuffer);
+
+            List<IRAttr*> offsetAttrs;
+            if (findNonuniformIndexInst(oldOffset))
+            {
+                offsetAttrs.add(getBuilder()->getAttr(kIROp_NonUniformAttr));
+            }
+            ioInfo.key.vals.add(
+                getBuilder()->getAttributedType(oldOffset->getDataType(), offsetAttrs));
+            ioInfo.newArgs.add(oldOffset);
+
+            ioInfo.key.vals.add(
+                getBuilder()->getAttributedType(oldAlignment->getDataType(), List<IRAttr*>()));
+            ioInfo.newArgs.add(oldAlignment);
+        }
         else if (isFieldAccessInst(oldArg))
         {
             // This is the case where the `oldArg` is
@@ -790,7 +812,6 @@ struct FunctionParameterSpecializationContext
         case kIROp_GetElement:
         case kIROp_RWStructuredBufferGetElementPtr:
         case kIROp_StructuredBufferLoad:
-        case kIROp_ByteAddressBufferLoad:
             return true;
         }
         return false;
@@ -898,6 +919,29 @@ struct FunctionParameterSpecializationContext
 
             return newVal;
         }
+        else if (oldArg->getOp() == kIROp_ByteAddressBufferLoad)
+        {
+            // Parallel to the element-access case above, with an extra `alignment` operand.
+            auto oldBuffer = oldArg->getOperand(0);
+            auto oldOffset = oldArg->getOperand(1);
+            auto oldAlignment = oldArg->getOperand(2);
+
+            auto newBuffer = getSpecializedValueForArg(ioInfo, oldBuffer);
+
+            auto builder = getBuilder();
+            auto newOffset = builder->createParam(oldOffset->getFullType());
+            ioInfo.newParams.add(newOffset);
+
+            auto newAlignment = builder->createParam(oldAlignment->getFullType());
+            ioInfo.newParams.add(newAlignment);
+
+            builder->setInsertInto(ioInfo.newBodyInsts);
+            IRInst* newOperands[] = {newBuffer, newOffset, newAlignment};
+            auto newVal =
+                builder->emitIntrinsicInst(oldArg->getFullType(), oldArg->getOp(), 3, newOperands);
+
+            return newVal;
+        }
         else if (isFieldAccessInst(oldArg))
         {
             // This is the case where the argument is
diff --git a/tests/optimization/buffer-load-specialize-byte-address.slang b/tests/optimization/buffer-load-specialize-byte-address.slang
new file mode 100644
index 00000000000..703f9e71117
--- /dev/null
+++ b/tests/optimization/buffer-load-specialize-byte-address.slang
@@ -0,0 +1,37 @@
+//TEST:SIMPLE(filecheck=SPV): -target spirv -O0 -fvk-use-scalar-layout
+
+// Specializing a callee whose argument is `ByteAddressBuffer.Load<T>()` must preserve
+// the `alignment` operand of `byteAddressBufferLoad`; check that `makeRay` is specialized
+// to take the offset and alignment as separate `uint` parameters.
+
+struct Frame
+{
+    float4x4 a;
+    float4x4 b;
+    float3 v;
+    float pad;
+};
+
+float3 makeRay(Frame f)
+{
+    float4 w = mul(float4(0, 0, 0, 1), f.b);
+    return w.xyz + f.v;
+}
+
+uniform ByteAddressBuffer buf;
+uniform RWTexture2D<float4> outTexture;
+uniform uint width;
+uniform uint height;
+
+[shader("compute")]
+[numthreads(8, 8, 1)]
+void computeMain(uint3 dtid : SV_DispatchThreadID)
+{
+    if (dtid.x >= width || dtid.y >= height) return;
+    Frame f = buf.Load<Frame>(0);
+    outTexture[dtid.xy] = float4(makeRay(f), 1);
+}
+
+// SPV: %makeRay = OpFunction %v3float None %{{.*}}
+// SPV-NEXT: %{{.*}} = OpFunctionParameter %uint
+// SPV-NEXT: %{{.*}} = OpFunctionParameter %uint