[llvm] [DirectX] Handle dx.RawBuffer in DXILResourceAccess (PR #121725)

[Justin Bogner via llvm-commits]

================
@@ -21,99 +21,206 @@
 
 using namespace llvm;
 
-static void replaceTypedBufferAccess(IntrinsicInst *II,
-                                     dxil::ResourceTypeInfo &RTI) {
-  const DataLayout &DL = II->getDataLayout();
+static Value *calculateGEPOffset(GetElementPtrInst *GEP, Value *PrevOffset,
+                                 dxil::ResourceTypeInfo &RTI) {
+  assert(!PrevOffset && "Non-constant GEP chains not handled yet");
+
+  const DataLayout &DL = GEP->getDataLayout();
+
+  uint64_t ScalarSize = 1;
+  if (RTI.isTyped()) {
+    Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
+    // We need the size of an element in bytes so that we can calculate the
+    // offset in elements given a total offset in bytes.
+    Type *ScalarType = ContainedType->getScalarType();
+    ScalarSize = DL.getTypeSizeInBits(ScalarType) / 8;
+  }
+
+  APInt ConstantOffset(DL.getIndexTypeSizeInBits(GEP->getType()), 0);
+  if (GEP->accumulateConstantOffset(DL, ConstantOffset)) {
+    APInt Scaled = ConstantOffset.udiv(ScalarSize);
+    return ConstantInt::get(Type::getInt32Ty(GEP->getContext()), Scaled);
+  }
+
+  auto IndexIt = GEP->idx_begin();
+  assert(cast<ConstantInt>(IndexIt)->getZExtValue() == 0 &&
+         "GEP is not indexing through pointer");
+  ++IndexIt;
+  Value *Offset = *IndexIt;
+  assert(++IndexIt == GEP->idx_end() && "Too many indices in GEP");
+  return Offset;
+}
+
+static void createTypedBufferStore(IntrinsicInst *II, StoreInst *SI,
+                                   Value *Offset, dxil::ResourceTypeInfo &RTI) {
+  IRBuilder<> Builder(SI);
+  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
+  Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());
+
+  Value *V = SI->getValueOperand();
+  if (V->getType() == ContainedType) {
+    // V is already the right type.
+    assert(!Offset && "store of whole element has offset?");
+  } else if (V->getType() == ContainedType->getScalarType()) {
+    // We're storing a scalar, so we need to load the current value and only
+    // replace the relevant part.
+    auto *Load = Builder.CreateIntrinsic(
+        LoadType, Intrinsic::dx_resource_load_typedbuffer,
+        {II->getOperand(0), II->getOperand(1)});
+    auto *Struct = Builder.CreateExtractValue(Load, {0});
+
+    // If we have an offset from seeing a GEP earlier, use that. Otherwise, 0.
+    if (!Offset)
+      Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
+    V = Builder.CreateInsertElement(Struct, V, Offset);
+  } else {
+    llvm_unreachable("Store to typed resource has invalid type");
+  }
+
+  auto *Inst = Builder.CreateIntrinsic(
+      Builder.getVoidTy(), Intrinsic::dx_resource_store_typedbuffer,
+      {II->getOperand(0), II->getOperand(1), V});
+  SI->replaceAllUsesWith(Inst);
+}
+
+static void createRawStore(IntrinsicInst *II, StoreInst *SI, Value *Offset) {
+  IRBuilder<> Builder(SI);
+
+  if (!Offset)
+    Offset = ConstantInt::get(Builder.getInt32Ty(), 0);
+  Value *V = SI->getValueOperand();
+  // TODO: break up larger types
+  auto *Inst = Builder.CreateIntrinsic(
+      Builder.getVoidTy(), Intrinsic::dx_resource_store_rawbuffer,
+      {II->getOperand(0), II->getOperand(1), Offset, V});
+  SI->replaceAllUsesWith(Inst);
+}
+
+static void createStoreIntrinsic(IntrinsicInst *II, StoreInst *SI,
+                                 Value *Offset, dxil::ResourceTypeInfo &RTI) {
+  switch (RTI.getResourceKind()) {
+  case dxil::ResourceKind::TypedBuffer:
+    return createTypedBufferStore(II, SI, Offset, RTI);
+  case dxil::ResourceKind::RawBuffer:
+  case dxil::ResourceKind::StructuredBuffer:
+    return createRawStore(II, SI, Offset);
+  case dxil::ResourceKind::Texture1D:
+  case dxil::ResourceKind::Texture2D:
+  case dxil::ResourceKind::Texture2DMS:
+  case dxil::ResourceKind::Texture3D:
+  case dxil::ResourceKind::TextureCube:
+  case dxil::ResourceKind::Texture1DArray:
+  case dxil::ResourceKind::Texture2DArray:
+  case dxil::ResourceKind::Texture2DMSArray:
+  case dxil::ResourceKind::TextureCubeArray:
+  case dxil::ResourceKind::FeedbackTexture2D:
+  case dxil::ResourceKind::FeedbackTexture2DArray:
+    // TODO: handle these
+    return;
+  case dxil::ResourceKind::CBuffer:
+  case dxil::ResourceKind::Sampler:
+  case dxil::ResourceKind::TBuffer:
+  case dxil::ResourceKind::RTAccelerationStructure:
+  case dxil::ResourceKind::Invalid:
+  case dxil::ResourceKind::NumEntries:
+    llvm_unreachable("Invalid resource kind for store");
+  }
+  llvm_unreachable("Unhandled case in switch");
+}
+
+static void createTypedBufferLoad(IntrinsicInst *II, LoadInst *LI,
+                                  Value *Offset, dxil::ResourceTypeInfo &RTI) {
+  IRBuilder<> Builder(LI);
+  Type *ContainedType = RTI.getHandleTy()->getTypeParameter(0);
+  Type *LoadType = StructType::get(ContainedType, Builder.getInt1Ty());
 
-  auto *HandleType = cast<TargetExtType>(II->getOperand(0)->getType());
-  assert(HandleType->getName() == "dx.TypedBuffer" &&
-         "Unexpected typed buffer type");
-  Type *ContainedType = HandleType->getTypeParameter(0);
+  Value *V =
+      Builder.CreateIntrinsic(LoadType, Intrinsic::dx_resource_load_typedbuffer,
+                              {II->getOperand(0), II->getOperand(1)});
+  V = Builder.CreateExtractValue(V, {0});
 
-  Type *LoadType =
-      StructType::get(ContainedType, Type::getInt1Ty(II->getContext()));
+  if (Offset)
+    V = Builder.CreateExtractElement(V, Offset);
 
-  // We need the size of an element in bytes so that we can calculate the offset
-  // in elements given a total offset in bytes later.
-  Type *ScalarType = ContainedType->getScalarType();
-  uint64_t ScalarSize = DL.getTypeSizeInBits(ScalarType) / 8;
+  LI->replaceAllUsesWith(V);
+}
 
+static void createRawLoad(IntrinsicInst *II, LoadInst *LI, Value *Offset) {
+  IRBuilder<> Builder(LI);
+  // TODO: break up larger types
----------------
bogner wrote:

I tried a few ways but the assert here is a little tricky to formulate, and we don't really end up with invalid DXIL anyway since the lowering of the load will fail if we get anything really silly here. I think I'll leave it as is for now.

https://github.com/llvm/llvm-project/pull/121725