[DirectX] Disallow ElementIndex for raw buffer accesses

llvm/docs/DirectX/DXILResources.rst

@@ -402,6 +402,11 @@ which matches DXIL. Unlike in the `RawBufferLoad`_ operation, we do not need
 arguments for the mask/type size and alignment, since we can calculate these
 from the return type of the load during lowering.
 
+Note that RawBuffer loads represent either "structured" accesses, as in HLSL's
+StructuredBuffer<T>, or a "raw" access, as in HLSL's "ByteAddressBuffer". The
+`%offset` parameter is only used for structured accesses, and *must* be
+`poison` for raw accesses.
+
 .. _RawBufferLoad: https://github.com/microsoft/DirectXShaderCompiler/blob/main/docs/DXIL.rst#rawbufferload
 
 .. list-table:: ``@llvm.dx.resource.load.rawbuffer``
@@ -442,7 +447,7 @@ Examples:
        @llvm.dx.resource.load.rawbuffer.f32.tdx.RawBuffer_i8_0_0_0t(
            target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
            i32 %byte_offset,
-           i32 0)
+           i32 poison)
 
    ; float4
    %ret = call {<4 x float>, i1}
@@ -454,7 +459,7 @@ Examples:
        @llvm.dx.resource.load.rawbuffer.v4f32.tdx.RawBuffer_i8_0_0_0t(
            target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
            i32 %byte_offset,
-           i32 0)
+           i32 poison)
 
    ; struct S0 { float4 f; int4 i; };
    %ret = call {<4 x float>, i1}
@@ -488,7 +493,7 @@ Examples:
        @llvm.dx.resource.load.rawbuffer.v4i64.tdx.RawBuffer_i8_0_0t(
            target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
            i32 %byte_offset,
-           i32 0)
+           i32 poison)
 
 Stores
 ------

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -272,6 +272,21 @@ class OpLowerer {
     return false;
   }
 
+  Error validateRawBufferElementIndex(Value *Resource, Value *ElementIndex) {
+    bool IsStruct = cast<RawBufferExtType>(Resource->getType())->isStructured();
+    bool IsPoison = isa<PoisonValue>(ElementIndex);
+
+    if (IsStruct && IsPoison)
+      return make_error<StringError>(
+          "Element index of structured buffer may not be poison",
+          inconvertibleErrorCode());
+    else if (!IsStruct && !IsPoison)
+      return make_error<StringError>(
+          "Element index of raw buffer must be poison",
+          inconvertibleErrorCode());
+    return Error::success();
+  }
+
   [[nodiscard]] bool lowerToCreateHandle(Function &F) {
     IRBuilder<> &IRB = OpBuilder.getIRB();
     Type *Int8Ty = IRB.getInt8Ty();
@@ -560,6 +575,11 @@ class OpLowerer {
       Value *Align =
           ConstantInt::get(Int32Ty, DL.getPrefTypeAlign(ScalarTy).value());
 
+      if (Error E = validateRawBufferElementIndex(CI->getOperand(0), Index1))
+        return E;
+      if (isa<PoisonValue>(Index1))
+        Index1 = UndefValue::get(Index1->getType());
+
       Expected<CallInst *> OpCall =
           MMDI.DXILVersion >= VersionTuple(1, 2)
               ? OpBuilder.tryCreateOp(OpCode::RawBufferLoad,
@@ -671,6 +691,13 @@ class OpLowerer {
       Value *Index0 = CI->getArgOperand(1);
       Value *Index1 = IsRaw ? CI->getArgOperand(2) : UndefValue::get(Int32Ty);
 
+      if (IsRaw) {
+        if (Error E = validateRawBufferElementIndex(CI->getOperand(0), Index1))
+          return E;
+        if (isa<PoisonValue>(Index1))
+          Index1 = UndefValue::get(Index1->getType());
+      }
+
       Value *Data = CI->getArgOperand(IsRaw ? 3 : 2);
       Type *DataTy = Data->getType();
       Type *ScalarTy = DataTy->getScalarType();

llvm/lib/Target/DirectX/DXILResourceAccess.cpp

@@ -120,19 +120,33 @@ static void createTypedBufferStore(IntrinsicInst *II, StoreInst *SI,
   SI->replaceAllUsesWith(Inst);
 }
 
-static void createRawStore(IntrinsicInst *II, StoreInst *SI) {
+static void createRawStore(IntrinsicInst *II, StoreInst *SI,
+                           dxil::ResourceTypeInfo &RTI) {
   const DataLayout &DL = SI->getDataLayout();
   IRBuilder<> Builder(SI);
 
   Value *V = SI->getValueOperand();
+  assert(!V->getType()->isAggregateType() &&
+         "Resource store should be scalar or vector type");
+
+  Value *Index = II->getOperand(1);
   // The offset for the rawbuffer load and store ops is always in bytes.
   uint64_t AccessSize = 1;
   Value *Offset =
       traverseGEPOffsets(DL, Builder, SI->getPointerOperand(), AccessSize);
-  // TODO: break up larger types
-  auto *Inst = Builder.CreateIntrinsic(
-      Builder.getVoidTy(), Intrinsic::dx_resource_store_rawbuffer,
-      {II->getOperand(0), II->getOperand(1), Offset, V});
+
+  // For raw buffer (ie, HLSL's ByteAddressBuffer), we need to fold the access
+  // entirely into the index.
+  if (!RTI.isStruct()) {
+    auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
+    if (!ConstantOffset || !ConstantOffset->isZero())
+      Index = Builder.CreateAdd(Index, Offset);
+    Offset = llvm::PoisonValue::get(Builder.getInt32Ty());
+  }
+
+  auto *Inst = Builder.CreateIntrinsic(Builder.getVoidTy(),
+                                       Intrinsic::dx_resource_store_rawbuffer,
+                                       {II->getOperand(0), Index, Offset, V});
   SI->replaceAllUsesWith(Inst);
 }
 
@@ -143,7 +157,7 @@ static void createStoreIntrinsic(IntrinsicInst *II, StoreInst *SI,
     return createTypedBufferStore(II, SI, RTI);
   case dxil::ResourceKind::RawBuffer:
   case dxil::ResourceKind::StructuredBuffer:
-    return createRawStore(II, SI);
+    return createRawStore(II, SI, RTI);
   case dxil::ResourceKind::Texture1D:
   case dxil::ResourceKind::Texture2D:
   case dxil::ResourceKind::Texture2DMS:
@@ -198,20 +212,33 @@ static void createTypedBufferLoad(IntrinsicInst *II, LoadInst *LI,
   LI->replaceAllUsesWith(V);
 }
 
-static void createRawLoad(IntrinsicInst *II, LoadInst *LI) {
+static void createRawLoad(IntrinsicInst *II, LoadInst *LI,
+                          dxil::ResourceTypeInfo &RTI) {
   const DataLayout &DL = LI->getDataLayout();
   IRBuilder<> Builder(LI);
 
-  // TODO: break up larger types
   Type *LoadType = StructType::get(LI->getType(), Builder.getInt1Ty());
+  assert(!LI->getType()->isAggregateType() &&
+         "Resource load should be scalar or vector type");
 
+  Value *Index = II->getOperand(1);
   // The offset for the rawbuffer load and store ops is always in bytes.
   uint64_t AccessSize = 1;
   Value *Offset =
       traverseGEPOffsets(DL, Builder, LI->getPointerOperand(), AccessSize);
+
+  // For raw buffer (ie, HLSL's ByteAddressBuffer), we need to fold the access
+  // entirely into the index.
+  if (!RTI.isStruct()) {
+    auto *ConstantOffset = dyn_cast<ConstantInt>(Offset);
+    if (!ConstantOffset || !ConstantOffset->isZero())
+      Index = Builder.CreateAdd(Index, Offset);
+    Offset = llvm::PoisonValue::get(Builder.getInt32Ty());
+  }
+
   Value *V =
       Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_rawbuffer,
-                              {II->getOperand(0), II->getOperand(1), Offset});
+                              {II->getOperand(0), Index, Offset});
   V = Builder.CreateExtractValue(V, {0});
 
   LI->replaceAllUsesWith(V);
@@ -367,7 +394,7 @@ static void createLoadIntrinsic(IntrinsicInst *II, LoadInst *LI,
     return createTypedBufferLoad(II, LI, RTI);
   case dxil::ResourceKind::RawBuffer:
   case dxil::ResourceKind::StructuredBuffer:
-    return createRawLoad(II, LI);
+    return createRawLoad(II, LI, RTI);
   case dxil::ResourceKind::CBuffer:
     return createCBufferLoad(II, LI, RTI);
   case dxil::ResourceKind::Texture1D:

llvm/test/CodeGen/DirectX/BufferLoad-sm61.ll

@@ -25,12 +25,12 @@ define void @loadv4f32_byte(i32 %offset) {
       @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_0_0_0(
           i32 0, i32 0, i32 1, i32 0, ptr null)
 
-  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f32 @dx.op.bufferLoad.f32(i32 68, %dx.types.Handle %{{.*}}, i32 %offset, i32 0)
+  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f32 @dx.op.bufferLoad.f32(i32 68, %dx.types.Handle %{{.*}}, i32 %offset, i32 undef)
   %load = call {<4 x float>, i1}
       @llvm.dx.resource.load.rawbuffer.f32.tdx.RawBuffer_i8_0_0_0t(
           target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
           i32 %offset,
-          i32 0)
+          i32 poison)
 
   ret void
 }

llvm/test/CodeGen/DirectX/BufferStore-sm61.ll

@@ -23,10 +23,10 @@ define void @storef32_byte(i32 %offset, float %data) {
       @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_0_0_0(
           i32 0, i32 0, i32 1, i32 0, ptr null)
 
-  ; CHECK: call void @dx.op.bufferStore.f32(i32 69, %dx.types.Handle %{{.*}}, i32 %offset, i32 0, float %data, float undef, float undef, float undef, i8 1)
+  ; CHECK: call void @dx.op.bufferStore.f32(i32 69, %dx.types.Handle %{{.*}}, i32 %offset, i32 undef, float %data, float undef, float undef, float undef, i8 1)
   call void @llvm.dx.resource.store.rawbuffer.f32(
       target("dx.RawBuffer", i8, 1, 0, 0) %buffer,
-      i32 %offset, i32 0, float %data)
+      i32 %offset, i32 poison, float %data)
 
   ret void
 }
@@ -59,10 +59,10 @@ define void @storev4f32_byte(i32 %offset, <4 x float> %data) {
   ; CHECK: [[DATA1:%.*]] = extractelement <4 x float> %data, i32 1
   ; CHECK: [[DATA2:%.*]] = extractelement <4 x float> %data, i32 2
   ; CHECK: [[DATA3:%.*]] = extractelement <4 x float> %data, i32 3
-  ; CHECK: call void @dx.op.bufferStore.f32(i32 69, %dx.types.Handle %{{.*}}, i32 %offset, i32 0, float [[DATA0]], float [[DATA1]], float [[DATA2]], float [[DATA3]], i8 15)
+  ; CHECK: call void @dx.op.bufferStore.f32(i32 69, %dx.types.Handle %{{.*}}, i32 %offset, i32 undef, float [[DATA0]], float [[DATA1]], float [[DATA2]], float [[DATA3]], i8 15)
   call void @llvm.dx.resource.store.rawbuffer.v4f32(
       target("dx.RawBuffer", i8, 1, 0, 0) %buffer,
-      i32 %offset, i32 0, <4 x float> %data)
+      i32 %offset, i32 poison, <4 x float> %data)
 
   ret void
 }

llvm/test/CodeGen/DirectX/RawBuffer-errors.ll

@@ -0,0 +1,78 @@
+; We use llc for this test so that we don't abort after the first error.
+; RUN: not llc %s -o /dev/null 2>&1 | FileCheck %s
+
+target triple = "dxil-pc-shadermodel6.6-compute"
+
+declare void @f32_user(float)
+
+; CHECK: error:
+; CHECK-SAME: in function loadrawzero
+; CHECK-SAME: Element index of raw buffer must be poison
+define void @loadrawzero(i32 %offset) "hlsl.export" {
+  %buffer = call target("dx.RawBuffer", i8, 0, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  %load = call {float, i1}
+      @llvm.dx.resource.load.rawbuffer(
+          target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
+          i32 %offset,
+          i32 0)
+  %data = extractvalue {float, i1} %load, 0
+
+  call void @f32_user(float %data)
+
+  ret void
+}
+
+; CHECK: error:
+; CHECK-SAME: in function loadstructundef
+; CHECK-SAME: Element index of structured buffer may not be poison
+define void @loadstructundef(i32 %index) "hlsl.export" {
+  %buffer = call target("dx.RawBuffer", float, 0, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  %load = call {float, i1}
+      @llvm.dx.resource.load.rawbuffer(
+          target("dx.RawBuffer", float, 0, 0, 0) %buffer,
+          i32 %index,
+          i32 poison)
+  %data = extractvalue {float, i1} %load, 0
+  call void @f32_user(float %data)
+
+  ret void
+}
+
+; CHECK: error:
+; CHECK-SAME: in function storerawzero
+; CHECK-SAME: Element index of raw buffer must be poison
+define void @storerawzero(i32 %offset, float %data) {
+  %buffer = call target("dx.RawBuffer", i8, 1, 0, 0)
+      @llvm.dx.resource.handlefrombinding(
+          i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  call void @llvm.dx.resource.store.rawbuffer(
+      target("dx.RawBuffer", i8, 1, 0, 0) %buffer,
+      i32 %offset, i32 0, float %data)
+
+  ret void
+}
+
+; CHECK: error:
+; CHECK-SAME: in function storestructundef
+; CHECK-SAME: Element index of structured buffer may not be poison
+define void @storestructundef(i32 %index, float %data) {
+  %buffer = call target("dx.RawBuffer", float, 1, 0, 0)
+      @llvm.dx.resource.handlefrombinding(
+          i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  call void @llvm.dx.resource.store.rawbuffer(
+      target("dx.RawBuffer", float, 1, 0, 0) %buffer,
+      i32 %index, i32 poison, float %data)
+
+  ret void
+}
+
+declare { float, i1 } @llvm.dx.resource.load.rawbuffer.f32.tdx.RawBuffer_i8_0_0_0t(target("dx.RawBuffer", i8, 0, 0, 0), i32, i32)
+declare { float, i1 } @llvm.dx.resource.load.rawbuffer.f32.tdx.RawBuffer_f32_0_0_0t(target("dx.RawBuffer", float, 0, 0, 0), i32, i32)
+declare void @llvm.dx.resource.store.rawbuffer.tdx.RawBuffer_i8_1_0_0t.f32(target("dx.RawBuffer", i8, 1, 0, 0), i32, i32, float)
+declare void @llvm.dx.resource.store.rawbuffer.tdx.RawBuffer_f32_1_0_0t.f32(target("dx.RawBuffer", float, 1, 0, 0), i32, i32, float)

llvm/test/CodeGen/DirectX/RawBufferLoad.ll

@@ -36,12 +36,12 @@ define void @loadf32_byte(i32 %offset) {
       @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_0_0_0(
           i32 0, i32 0, i32 1, i32 0, ptr null)
 
-  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f32 @dx.op.rawBufferLoad.f32(i32 139, %dx.types.Handle %{{.*}}, i32 %offset, i32 0, i8 1, i32 4)
+  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f32 @dx.op.rawBufferLoad.f32(i32 139, %dx.types.Handle %{{.*}}, i32 %offset, i32 undef, i8 1, i32 4)
   %load = call {float, i1}
       @llvm.dx.resource.load.rawbuffer.f32.tdx.RawBuffer_i8_0_0_0t(
           target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
           i32 %offset,
-          i32 0)
+          i32 poison)
   %data = extractvalue {float, i1} %load, 0
 
   ; CHECK: [[VAL:%.*]] = extractvalue %dx.types.ResRet.f32 [[DATA]], 0
@@ -85,12 +85,12 @@ define void @loadv4f32_byte(i32 %offset) {
       @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_0_0_0(
           i32 0, i32 0, i32 1, i32 0, ptr null)
 
-  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f32 @dx.op.rawBufferLoad.f32(i32 139, %dx.types.Handle %{{.*}}, i32 %offset, i32 0, i8 15, i32 4)
+  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f32 @dx.op.rawBufferLoad.f32(i32 139, %dx.types.Handle %{{.*}}, i32 %offset, i32 undef, i8 15, i32 4)
   %load = call {<4 x float>, i1}
       @llvm.dx.resource.load.rawbuffer.f32.tdx.RawBuffer_i8_0_0_0t(
           target("dx.RawBuffer", i8, 0, 0, 0) %buffer,
           i32 %offset,
-          i32 0)
+          i32 poison)
   %data = extractvalue {<4 x float>, i1} %load, 0
 
   ; CHECK: extractvalue %dx.types.ResRet.f32 [[DATA]], 0
@@ -212,9 +212,9 @@ define void @loadv4f64_byte(i32 %offset) {
       @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_0_0_0(
           i32 0, i32 0, i32 1, i32 0, ptr null)
 
-  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f64 @dx.op.rawBufferLoad.f64(i32 139, %dx.types.Handle %{{.*}}, i32 %offset, i32 0, i8 15, i32 8)
+  ; CHECK: [[DATA:%.*]] = call %dx.types.ResRet.f64 @dx.op.rawBufferLoad.f64(i32 139, %dx.types.Handle %{{.*}}, i32 %offset, i32 undef, i8 15, i32 8)
   %load = call {<4 x double>, i1} @llvm.dx.resource.load.rawbuffer.v4i64(
-      target("dx.RawBuffer", i8, 0, 0, 0) %buffer, i32 %offset, i32 0)
+      target("dx.RawBuffer", i8, 0, 0, 0) %buffer, i32 %offset, i32 poison)
   %data = extractvalue {<4 x double>, i1} %load, 0
 
   ; CHECK: extractvalue %dx.types.ResRet.f64 [[DATA]], 0

llvm/test/CodeGen/DirectX/RawBufferStore.ll

@@ -22,10 +22,10 @@ define void @storef32_byte(i32 %offset, float %data) {
       @llvm.dx.resource.handlefrombinding.tdx.RawBuffer_i8_0_0_0(
           i32 0, i32 0, i32 1, i32 0, ptr null)
 
-  ; CHECK: call void @dx.op.rawBufferStore.f32(i32 140, %dx.types.Handle %buffer_annot, i32 %offset, i32 0, float %data, float undef, float undef, float undef, i8 1, i32 4)
+  ; CHECK: call void @dx.op.rawBufferStore.f32(i32 140, %dx.types.Handle %buffer_annot, i32 %offset, i32 undef, float %data, float undef, float undef, float undef, i8 1, i32 4)
   call void @llvm.dx.resource.store.rawbuffer.f32(
       target("dx.RawBuffer", i8, 1, 0, 0) %buffer,
-      i32 %offset, i32 0, float %data)
+      i32 %offset, i32 poison, float %data)
 
   ret void
 }
@@ -58,10 +58,10 @@ define void @storev4f32_byte(i32 %offset, <4 x float> %data) {
   ; CHECK: [[DATA1:%.*]] = extractelement <4 x float> %data, i32 1
   ; CHECK: [[DATA2:%.*]] = extractelement <4 x float> %data, i32 2
   ; CHECK: [[DATA3:%.*]] = extractelement <4 x float> %data, i32 3
-  ; CHECK: call void @dx.op.rawBufferStore.f32(i32 140, %dx.types.Handle %buffer_annot, i32 %offset, i32 0, float [[DATA0]], float [[DATA1]], float [[DATA2]], float [[DATA3]], i8 15, i32 4)
+  ; CHECK: call void @dx.op.rawBufferStore.f32(i32 140, %dx.types.Handle %buffer_annot, i32 %offset, i32 undef, float [[DATA0]], float [[DATA1]], float [[DATA2]], float [[DATA3]], i8 15, i32 4)
   call void @llvm.dx.resource.store.rawbuffer.v4f32(
       target("dx.RawBuffer", i8, 1, 0, 0) %buffer,
-      i32 %offset, i32 0, <4 x float> %data)
+      i32 %offset, i32 poison, <4 x float> %data)
 
   ret void
 }
@@ -135,10 +135,10 @@ define void @storev4f64_byte(i32 %offset, <4 x double> %data) {
   ; CHECK: [[DATA1:%.*]] = extractelement <4 x double> %data, i32 1
   ; CHECK: [[DATA2:%.*]] = extractelement <4 x double> %data, i32 2
   ; CHECK: [[DATA3:%.*]] = extractelement <4 x double> %data, i32 3
-  ; CHECK: call void @dx.op.rawBufferStore.f64(i32 140, %dx.types.Handle %buffer_annot, i32 %offset, i32 0, double [[DATA0]], double [[DATA1]], double [[DATA2]], double [[DATA3]], i8 15, i32 8)
+  ; CHECK: call void @dx.op.rawBufferStore.f64(i32 140, %dx.types.Handle %buffer_annot, i32 %offset, i32 undef, double [[DATA0]], double [[DATA1]], double [[DATA2]], double [[DATA3]], i8 15, i32 8)
   call void @llvm.dx.resource.store.rawbuffer.v4i64(
       target("dx.RawBuffer", i8, 1, 0, 0) %buffer,
-      i32 %offset, i32 0, <4 x double> %data)
+      i32 %offset, i32 poison, <4 x double> %data)
 
   ret void
 }