Fix dropped alignment operand in byteAddressBufferLoad specialization

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

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