Lower native vector raw buffers load/stores into new ops

When the loaded/stored type is a vector of more than 1 element, the
shader model is 6.9 or higher, and the operation is on a raw buffer,
enable the generation of a native vector raw buffer load or store.

Incidental removal of unused parameter in load translation

add validation and compute shader tests

Author: Greg Roth

lib/HLSL/HLOperationLower.cpp

@@ -3953,6 +3953,11 @@ struct ResLoadHelper {
       : intrinsicOpCode(IntrinsicOp::Num_Intrinsics), handle(h), retVal(Inst),
         addr(idx), offset(Offset), status(nullptr), mipLevel(mip) {
     opcode = LoadOpFromResKind(RK);
+    Type *Ty = Inst->getType();
+    if (opcode == OP::OpCode::RawBufferLoad && Ty->isVectorTy() &&
+        Ty->getVectorNumElements() > 1 &&
+        Inst->getModule()->GetHLModule().GetShaderModel()->IsSM69Plus())
+      opcode = OP::OpCode::RawBufferVectorLoad;
   }
   OP::OpCode opcode;
   IntrinsicOp intrinsicOpCode;
@@ -4022,6 +4027,14 @@ ResLoadHelper::ResLoadHelper(CallInst *CI, DxilResource::Kind RK,
     if (RC == DxilResourceBase::Class::SRV)
       OffsetIdx = IsMS ? HLOperandIndex::kTex2DMSLoadOffsetOpIdx
                        : HLOperandIndex::kTexLoadOffsetOpIdx;
+  } else if (opcode == OP::OpCode::RawBufferLoad) {
+    // If native vectors are available and this load had a vector
+    // with more than one elements, convert the RawBufferLod to the
+    // native vector variant RawBufferVectorLoad.
+    Type *Ty = CI->getType();
+    if (Ty->isVectorTy() && Ty->getVectorNumElements() > 1 &&
+        CI->getModule()->GetHLModule().GetShaderModel()->IsSM69Plus())
+      opcode = OP::OpCode::RawBufferVectorLoad;
   }
 
   // Set offset.
@@ -4079,7 +4092,7 @@ Value *GenerateRawBufLd(Value *handle, Value *bufIdx, Value *offset,
 // Sets up arguments for buffer load call.
 static SmallVector<Value *, 10> GetBufLoadArgs(ResLoadHelper helper,
                                                HLResource::Kind RK,
-                                               IRBuilder<> Builder, Type *EltTy,
+                                               IRBuilder<> Builder,
                                                unsigned LdSize) {
   OP::OpCode opcode = helper.opcode;
   llvm::Constant *opArg = Builder.getInt32((uint32_t)opcode);
@@ -4127,6 +4140,7 @@ static SmallVector<Value *, 10> GetBufLoadArgs(ResLoadHelper helper,
     // If not TextureLoad, it could be a typed or raw buffer load.
     // They have mostly similar arguments.
     DXASSERT(opcode == OP::OpCode::RawBufferLoad ||
+                 opcode == OP::OpCode::RawBufferVectorLoad ||
                  opcode == OP::OpCode::BufferLoad,
              "Wrong opcode in get load args");
     Args.emplace_back(
@@ -4137,6 +4151,9 @@ static SmallVector<Value *, 10> GetBufLoadArgs(ResLoadHelper helper,
       // Unlike typed buffer load, raw buffer load has mask and alignment.
       Args.emplace_back(nullptr);      // Mask will be added later %4.
       Args.emplace_back(alignmentVal); // alignment @5.
+    } else if (opcode == OP::OpCode::RawBufferVectorLoad) {
+      // RawBufferVectorLoad takes just alignment, no mask.
+      Args.emplace_back(alignmentVal); // alignment @4
     }
   }
   return Args;
@@ -4162,18 +4179,19 @@ Value *TranslateBufLoad(ResLoadHelper &helper, HLResource::Kind RK,
   if (isBool || (is64 && isTyped))
     EltTy = Builder.getInt32Ty();
 
-  // 64-bit types are stored as int32 pairs in typed buffers.
+  // Adjust number of components as needed.
   if (is64 && isTyped) {
+    // 64-bit types are stored as int32 pairs in typed buffers.
     DXASSERT(NumComponents <= 2, "Typed buffers only allow 4 dwords.");
     NumComponents *= 2;
+  } else if (opcode == OP::OpCode::RawBufferVectorLoad) {
+    // Native vector loads only have a single vector element in ResRet.
+    EltTy = VectorType::get(EltTy, NumComponents);
+    NumComponents = 1;
   }
 
   unsigned LdSize = DL.getTypeAllocSize(EltTy);
-
-  SmallVector<Value *, 4> Elts(NumComponents);
-
-  SmallVector<Value *, 10> Args =
-      GetBufLoadArgs(helper, RK, Builder, EltTy, LdSize);
+  SmallVector<Value *, 10> Args = GetBufLoadArgs(helper, RK, Builder, LdSize);
 
   // Keep track of the first load for debug info migration.
   Value *FirstLd = nullptr;
@@ -4185,9 +4203,10 @@ Value *TranslateBufLoad(ResLoadHelper &helper, HLResource::Kind RK,
   else if (RK == DxilResource::Kind::StructuredBuffer)
     OffsetIdx = DXIL::OperandIndex::kRawBufferLoadElementOffsetOpIdx;
 
-  // Create calls to function object.
+  // Create call(s) to function object and collect results in Elts.
   // Typed buffer loads are limited to one load of up to 4 32-bit values.
   // Raw buffer loads might need multiple loads in chunks of 4.
+  SmallVector<Value *, 4> Elts(NumComponents);
   for (unsigned i = 0; i < NumComponents;) {
     // Load 4 elements or however many less than 4 are left to load.
     unsigned chunkSize = std::min(NumComponents - i, 4U);
@@ -4197,7 +4216,7 @@ Value *TranslateBufLoad(ResLoadHelper &helper, HLResource::Kind RK,
       Args[DXIL::OperandIndex::kRawBufferLoadMaskOpIdx] =
           GetRawBufferMaskForETy(EltTy, chunkSize, OP);
       // If we've loaded a chunk already, update offset to next chunk.
-      if (FirstLd != nullptr && opcode == OP::OpCode::RawBufferLoad)
+      if (FirstLd != nullptr)
         Args[OffsetIdx] =
             Builder.CreateAdd(Args[OffsetIdx], OP->GetU32Const(4 * LdSize));
     }
@@ -4206,8 +4225,13 @@ Value *TranslateBufLoad(ResLoadHelper &helper, HLResource::Kind RK,
     Value *Ld = Builder.CreateCall(F, Args, OP::GetOpCodeName(opcode));
 
     // Extract elements from returned ResRet.
-    for (unsigned j = 0; j < chunkSize; j++, i++)
-      Elts[i] = Builder.CreateExtractValue(Ld, j);
+    // Native vector loads just have one vector element in the ResRet.
+    // Others have up to four scalars that need to be individually extracted.
+    if (opcode == OP::OpCode::RawBufferVectorLoad)
+      Elts[i++] = Builder.CreateExtractValue(Ld, 0);
+    else
+      for (unsigned j = 0; j < chunkSize; j++, i++)
+        Elts[i] = Builder.CreateExtractValue(Ld, j);
 
     // Update status.
     UpdateStatus(Ld, helper.status, Builder, OP);
@@ -4245,9 +4269,10 @@ Value *TranslateBufLoad(ResLoadHelper &helper, HLResource::Kind RK,
     }
   }
 
-  // Package elements into a vector.
+  // Package elements into a vector as needed.
   Value *retValNew = nullptr;
-  if (!Ty->isVectorTy()) {
+  // Scalar or native vector loads need not construct vectors from elements.
+  if (!Ty->isVectorTy() || opcode == OP::OpCode::RawBufferVectorLoad) {
     retValNew = Elts[0];
   } else {
     retValNew = UndefValue::get(VectorType::get(EltTy, NumComponents));
@@ -4345,6 +4370,10 @@ void TranslateStore(DxilResource::Kind RK, Value *handle, Value *val,
   case DxilResource::Kind::StructuredBuffer:
     IsTyped = false;
     opcode = OP::OpCode::RawBufferStore;
+    // Where shader model and type allows, use vector store intrinsic.
+    if (OP->GetModule()->GetHLModule().GetShaderModel()->IsSM69Plus() &&
+        Ty->isVectorTy() && Ty->getVectorNumElements() > 1)
+      opcode = OP::OpCode::RawBufferVectorStore;
     break;
   case DxilResource::Kind::TypedBuffer:
     opcode = OP::OpCode::BufferStore;
@@ -4387,7 +4416,6 @@ void TranslateStore(DxilResource::Kind RK, Value *handle, Value *val,
     EltTy = i32Ty;
   }
 
-  Function *F = OP->GetOpFunc(opcode, EltTy);
   llvm::Constant *opArg = OP->GetU32Const((unsigned)opcode);
 
   llvm::Value *undefI =
@@ -4401,6 +4429,7 @@ void TranslateStore(DxilResource::Kind RK, Value *handle, Value *val,
 
   unsigned OffsetIdx = 0;
   if (opcode == OP::OpCode::RawBufferStore ||
+      opcode == OP::OpCode::RawBufferVectorStore ||
       opcode == OP::OpCode::BufferStore) {
     // Append Coord0 (Index) value.
     if (Idx->getType()->isVectorTy()) {
@@ -4420,7 +4449,6 @@ void TranslateStore(DxilResource::Kind RK, Value *handle, Value *val,
       OffsetIdx = storeArgs.size() - 1;
 
     // Coord1 (Offset).
-    // Only relevant when storing more than 4 elements to structured buffers.
     storeArgs.emplace_back(offset);
   } else {
     // texture store
@@ -4441,6 +4469,16 @@ void TranslateStore(DxilResource::Kind RK, Value *handle, Value *val,
     // TODO: support mip for texture ST
   }
 
+  // RawBufferVectorStore only takes a single value and alignment arguments.
+  if (opcode == DXIL::OpCode::RawBufferVectorStore) {
+    storeArgs.emplace_back(val);
+    storeArgs.emplace_back(Alignment);
+    Function *F = OP->GetOpFunc(DXIL::OpCode::RawBufferVectorStore, Ty);
+    Builder.CreateCall(F, storeArgs);
+    return;
+  }
+  Function *F = OP->GetOpFunc(opcode, EltTy);
+
   constexpr unsigned MaxStoreElemCount = 4;
   const unsigned CompCount = Ty->isVectorTy() ? Ty->getVectorNumElements() : 1;
   const unsigned StoreInstCount =

tools/clang/test/CodeGenDXIL/hlsl/intrinsics/buffer-load-stores-sm69.hlsl

@@ -0,0 +1,91 @@
+// RUN: %dxc -DTYPE=float    -DNUM=4 -T vs_6_9 %s | FileCheck %s
+// RUN: %dxc -DTYPE=bool     -DNUM=4 -T vs_6_9 %s | FileCheck %s --check-prefixes=CHECK,I1
+// RUN: %dxc -DTYPE=uint64_t -DNUM=2 -T vs_6_9 %s | FileCheck %s
+// RUN: %dxc -DTYPE=double   -DNUM=2 -T vs_6_9 %s | FileCheck %s
+
+// RUN: %dxc -DTYPE=float    -DNUM=6  -T vs_6_9 %s | FileCheck %s
+// RUN: %dxc -DTYPE=bool     -DNUM=13 -T vs_6_9 %s | FileCheck %s --check-prefixes=CHECK,I1
+// RUN: %dxc -DTYPE=uint64_t -DNUM=24 -T vs_6_9 %s | FileCheck %s
+// RUN: %dxc -DTYPE=double   -DNUM=32 -T vs_6_9 %s | FileCheck %s
+
+///////////////////////////////////////////////////////////////////////
+// Test codegen for various load and store operations and conversions
+//  for different scalar/vector buffer types and indices.
+///////////////////////////////////////////////////////////////////////
+
+// CHECK: %dx.types.ResRet.[[VTY:v[0-9]*[a-z][0-9][0-9]]] = type { <[[NUM:[0-9]*]] x [[TYPE:[a-z_0-9]*]]>, i32 }
+
+  ByteAddressBuffer RoByBuf : register(t1);
+RWByteAddressBuffer RwByBuf : register(u1);
+
+StructuredBuffer< vector<TYPE, NUM> > RoStBuf : register(t2);
+RWStructuredBuffer< vector<TYPE, NUM>  > RwStBuf : register(u2);
+
+ConsumeStructuredBuffer<vector<TYPE, NUM> > CnStBuf : register(u4);
+AppendStructuredBuffer<vector<TYPE, NUM> > ApStBuf  : register(u5);
+
+// CHECK-LABEL: define void @main
+[shader("vertex")]
+void main(uint ix[2] : IX) {
+  // ByteAddressBuffer Tests
+
+  // CHECK-DAG: [[HDLROBY:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, %dx.types.ResBind { i32 1, i32 1, i32 0, i8 0 }, i32 1, i1 false)
+  // CHECK-DAG: [[HDLRWBY:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, %dx.types.ResBind { i32 1, i32 1, i32 0, i8 1 }, i32 1, i1 false)
+
+  // CHECK-DAG: [[HDLROST:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, %dx.types.ResBind { i32 2, i32 2, i32 0, i8 0 }, i32 2, i1 false)
+  // CHECK-DAG: [[HDLRWST:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, %dx.types.ResBind { i32 2, i32 2, i32 0, i8 1 }, i32 2, i1 false)
+
+  // CHECK-DAG: [[HDLCON:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, %dx.types.ResBind { i32 4, i32 4, i32 0, i8 1 }, i32 4, i1 false)
+  // CHECK-DAG: [[HDLAPP:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, %dx.types.ResBind { i32 5, i32 5, i32 0, i8 1 }, i32 5, i1 false)
+
+  // CHECK: [[IX0:%.*]] = call i32 @dx.op.loadInput.i32(i32 4,
+
+  // CHECK: [[ANHDLRWBY:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HDLRWBY]]
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLRWBY]], i32 [[IX0]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  babElt1 = RwByBuf.Load< vector<TYPE, NUM>  >(ix[0]);
+
+  // CHECK: [[ANHDLROBY:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HDLROBY]]
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLROBY]], i32 [[IX0]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  babElt2 = RoByBuf.Load< vector<TYPE, NUM>  >(ix[0]);
+
+  // I1: zext <[[NUM]] x i1> %{{.*}} to <[[NUM]] x i32>
+  // CHECK: all void @dx.op.rawBufferVectorStore.[[VTY]](i32 304, %dx.types.Handle [[ANHDLRWBY]], i32 [[IX0]]
+  RwByBuf.Store< vector<TYPE, NUM>  >(ix[0], babElt1 + babElt2);
+
+  // StructuredBuffer Tests
+  // CHECK: [[ANHDLRWST:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HDLRWST]]
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLRWST]], i32 [[IX0]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  stbElt1 = RwStBuf.Load(ix[0]);
+  // CHECK: [[IX1:%.*]] = call i32 @dx.op.loadInput.i32(i32 4,
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLRWST]], i32 [[IX1]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  stbElt2 = RwStBuf[ix[1]];
+
+  // CHECK: [[ANHDLROST:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HDLROST]]
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLROST]], i32 [[IX0]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  stbElt3 = RoStBuf.Load(ix[0]);
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLROST]], i32 [[IX1]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  stbElt4 = RoStBuf[ix[1]];
+
+  // I1: zext <[[NUM]] x i1> %{{.*}} to <[[NUM]] x i32>
+  // CHECK: all void @dx.op.rawBufferVectorStore.[[VTY]](i32 304, %dx.types.Handle [[ANHDLRWST]], i32 [[IX0]]
+  RwStBuf[ix[0]] = stbElt1 + stbElt2 + stbElt3 + stbElt4;
+
+  // {Append/Consume}StructuredBuffer Tests
+  // CHECK: [[ANHDLCON:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HDLCON]]
+  // CHECK: [[CONIX:%.*]] = call i32 @dx.op.bufferUpdateCounter(i32 70, %dx.types.Handle [[ANHDLCON]], i8 -1) 
+  // CHECK: call %dx.types.ResRet.[[VTY]] @dx.op.rawBufferVectorLoad.[[VTY]](i32 303, %dx.types.Handle [[ANHDLCON]], i32 [[CONIX]]
+  // I1: icmp ne <[[NUM]] x i32> %{{.*}}, zeroinitializer
+  vector<TYPE, NUM>  cnElt = CnStBuf.Consume();
+
+  // CHECK: [[ANHDLAPP:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HDLAPP]]
+  // CHECK: [[APPIX:%.*]] = call i32 @dx.op.bufferUpdateCounter(i32 70, %dx.types.Handle [[ANHDLAPP]], i8 1) 
+  // I1: zext <[[NUM]] x i1> %{{.*}} to <[[NUM]] x i32>
+  // CHECK: all void @dx.op.rawBufferVectorStore.[[VTY]](i32 304, %dx.types.Handle [[ANHDLAPP]], i32 [[APPIX]]
+  ApStBuf.Append(cnElt);
+}