[DirectX] Add llvm.dx.resource.getdimensions intrinsic and lowering to DXIL

[hekota]

Introduces LLVM intrinsic lvm.dx.resource.getdimensions and its lowering to DXIL op op.dx.getDimensions.
The intrinsic will be used to implement various GetDimension overload on buffers and textures.

Closes #112982

[llvmbot]

@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-backend-directx

Author: Helena Kotas (hekota)

Changes

Introduces LLVM intrinsic lvm.dx.resource.getdimensions and its lowering to DXIL op op.dx.getDimensions.
The intrinsic will be used to implement various GetDimension overload on buffers and textures.

Closes #112982

---

Full diff: https://github.com/llvm/llvm-project/pull/161753.diff

5 Files Affected:

• (modified) llvm/include/llvm/IR/IntrinsicsDirectX.td (+4)
• (modified) llvm/lib/Target/DirectX/DXIL.td (+8)
• (modified) llvm/lib/Target/DirectX/DXILOpBuilder.cpp (+8)
• (modified) llvm/lib/Target/DirectX/DXILOpLowering.cpp (+27)
• (added) llvm/test/CodeGen/DirectX/bufferGetDimensions.ll (+34)

diff --git a/llvm/include/llvm/IR/IntrinsicsDirectX.td b/llvm/include/llvm/IR/IntrinsicsDirectX.td
index 570d6bc35cbd0..38cd3f3e07a87 100644
--- a/llvm/include/llvm/IR/IntrinsicsDirectX.td
+++ b/llvm/include/llvm/IR/IntrinsicsDirectX.td
@@ -77,6 +77,10 @@ def int_dx_resource_updatecounter
     : DefaultAttrsIntrinsic<[llvm_i32_ty], [llvm_any_ty, llvm_i8_ty],
                             [IntrInaccessibleMemOrArgMemOnly]>;
 
+def int_dx_resource_getdimensions
+    : DefaultAttrsIntrinsic<[llvm_i32_ty, llvm_i32_ty, llvm_i32_ty, llvm_i32_ty], [llvm_any_ty, llvm_i32_ty],
+                            [IntrReadMem]>;
+
 // Cast between target extension handle types and dxil-style opaque handles
 def int_dx_resource_casthandle : Intrinsic<[llvm_any_ty], [llvm_any_ty]>;
 
diff --git a/llvm/lib/Target/DirectX/DXIL.td b/llvm/lib/Target/DirectX/DXIL.td
index 228114c5c24b2..44c48305f2832 100644
--- a/llvm/lib/Target/DirectX/DXIL.td
+++ b/llvm/lib/Target/DirectX/DXIL.td
@@ -57,6 +57,7 @@ def ResBindTy : DXILOpParamType;
 def ResPropsTy : DXILOpParamType;
 def SplitDoubleTy : DXILOpParamType;
 def BinaryWithCarryTy : DXILOpParamType;
+def DimensionsTy : DXILOpParamType;
 
 class DXILOpClass;
 
@@ -901,6 +902,13 @@ def CheckAccessFullyMapped : DXILOp<71, checkAccessFullyMapped> {
   let attributes = [Attributes<DXIL1_0, [ReadOnly]>];
 }
 
+def GetDimensions : DXILOp<72, getDimensions> {
+  let Doc = "gets the dimensions of a buffer or texture";
+  let arguments = [HandleTy, Int32Ty];
+  let result = DimensionsTy;
+  let stages = [Stages<DXIL1_0, [all_stages]>];
+}
+
 def Barrier : DXILOp<80, barrier> {
   let Doc = "inserts a memory barrier in the shader";
   let intrinsics = [
diff --git a/llvm/lib/Target/DirectX/DXILOpBuilder.cpp b/llvm/lib/Target/DirectX/DXILOpBuilder.cpp
index 1aed8f9867231..944b2e6433988 100644
--- a/llvm/lib/Target/DirectX/DXILOpBuilder.cpp
+++ b/llvm/lib/Target/DirectX/DXILOpBuilder.cpp
@@ -261,6 +261,12 @@ static StructType *getBinaryWithCarryType(LLVMContext &Context) {
   return StructType::create({Int32Ty, Int1Ty}, "dx.types.i32c");
 }
 
+static StructType *getDimensionsType(LLVMContext &Ctx) {
+  Type *Int32Ty = Type::getInt32Ty(Ctx);
+  return getOrCreateStructType("dx.types.Dimensions",
+                               {Int32Ty, Int32Ty, Int32Ty, Int32Ty}, Ctx);
+}
+
 static Type *getTypeFromOpParamType(OpParamType Kind, LLVMContext &Ctx,
                                     Type *OverloadTy) {
   switch (Kind) {
@@ -318,6 +324,8 @@ static Type *getTypeFromOpParamType(OpParamType Kind, LLVMContext &Ctx,
     return getSplitDoubleType(Ctx);
   case OpParamType::BinaryWithCarryTy:
     return getBinaryWithCarryType(Ctx);
+  case OpParamType::DimensionsTy:
+    return getDimensionsType(Ctx);
   }
   llvm_unreachable("Invalid parameter kind");
   return nullptr;
diff --git a/llvm/lib/Target/DirectX/DXILOpLowering.cpp b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
index 610d8b63bba27..34bb450bce7d0 100644
--- a/llvm/lib/Target/DirectX/DXILOpLowering.cpp
+++ b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
@@ -627,6 +627,30 @@ class OpLowerer {
     });
   }
 
+  [[nodiscard]] bool lowerGetDimensions(Function &F) {
+    IRBuilder<> &IRB = OpBuilder.getIRB();
+    Type *Int32Ty = IRB.getInt32Ty();
+
+    return replaceFunction(F, [&](CallInst *CI) -> Error {
+      IRB.SetInsertPoint(CI);
+      Value *Handle =
+          createTmpHandleCast(CI->getArgOperand(0), OpBuilder.getHandleType());
+      Value *Op1 = CI->getArgOperand(1);
+      if (isa<llvm::PoisonValue>(Op1))
+        Op1 = UndefValue::get(Int32Ty);
+
+      Expected<CallInst *> OpCall = OpBuilder.tryCreateOp(
+          OpCode::GetDimensions, {Handle, Op1}, CI->getName(), Int32Ty);
+      if (Error E = OpCall.takeError())
+        return E;
+      if (Error E = replaceNamedStructUses(CI, *OpCall))
+        return E;
+
+      CI->eraseFromParent();
+      return Error::success();
+    });
+  }
+
   [[nodiscard]] bool lowerGetPointer(Function &F) {
     // These should have already been handled in DXILResourceAccess, so we can
     // just clean up the dead prototype.
@@ -934,6 +958,9 @@ class OpLowerer {
       case Intrinsic::dx_resource_updatecounter:
         HasErrors |= lowerUpdateCounter(F);
         break;
+      case Intrinsic::dx_resource_getdimensions:
+        HasErrors |= lowerGetDimensions(F);
+        break;
       case Intrinsic::ctpop:
         HasErrors |= lowerCtpopToCountBits(F);
         break;
diff --git a/llvm/test/CodeGen/DirectX/bufferGetDimensions.ll b/llvm/test/CodeGen/DirectX/bufferGetDimensions.ll
new file mode 100644
index 0000000000000..7805426181b35
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/bufferGetDimensions.ll
@@ -0,0 +1,34 @@
+; RUN: opt -S -dxil-op-lower %s | FileCheck %s
+
+target triple = "dxil-pc-shadermodel6.6-compute"
+
+define i32 @test_getdimensions_no_mips() {
+  ; CHECK: [[HANDLE1:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, 
+  ; CHECK-NEXT: [[ANNOT_HANDLE1:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HANDLE1]]
+  %handle1 = call target("dx.TypedBuffer", <4 x float>, 0, 0, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  ; CHECK-NEXT: [[RETVAL1:%.*]] = call %dx.types.Dimensions @dx.op.getDimensions(i32 72, %dx.types.Handle [[ANNOT_HANDLE1]], i32 undef)
+  %1 = call { i32, i32, i32, i32 } @llvm.dx.resource.getdimensions.tdx.RawBuffer_i32_1_0t(target("dx.TypedBuffer", <4 x float>, 0, 0, 0) %handle1, i32 poison)
+  
+  ; CHECK-NEXT: %[[DIM1:.*]] = extractvalue %dx.types.Dimensions [[RETVAL1]], 0
+  %2 = extractvalue { i32, i32, i32, i32 } %1, 0
+  
+  ; CHECK-NEXT: ret i32 %[[DIM1]]
+  ret i32 %2
+}
+
+
+define i32 @test_getdimensions_with_0_mips() {
+  ; CHECK: [[HANDLE2:%.*]] = call %dx.types.Handle @dx.op.createHandleFromBinding(i32 217, 
+  ; CHECK-NEXT: [[ANNOT_HANDLE2:%.*]] = call %dx.types.Handle @dx.op.annotateHandle(i32 216, %dx.types.Handle [[HANDLE2]]
+  %handle1 = call target("dx.RawBuffer", float, 0, 0) @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, ptr null)
+
+  ; CHECK-NEXT: [[RETVAL2:%.*]] = call %dx.types.Dimensions @dx.op.getDimensions(i32 72, %dx.types.Handle [[ANNOT_HANDLE2]], i32 0)
+  %1 = call { i32, i32, i32, i32 } @llvm.dx.resource.getdimensions.tdx.RawBuffer_i32_1_0t(target("dx.RawBuffer", float, 0, 0) %handle1, i32 0)
+  
+  ; CHECK-NEXT: %[[DIM2:.*]] = extractvalue %dx.types.Dimensions [[RETVAL2]], 0
+  %2 = extractvalue { i32, i32, i32, i32 } %1, 0
+  
+  ; CHECK-NEXT: ret i32 %[[DIM2]]
+  ret i32 %2
+}

[github-actions]

⚠️ undef deprecator found issues in your code. ⚠️

You can test this locally with the following command:

git diff -U0 --pickaxe-regex -S '([^a-zA-Z0-9#_-]undef[^a-zA-Z0-9_-]|UndefValue::get)' 'HEAD~1' HEAD llvm/test/CodeGen/DirectX/bufferGetDimensions.ll llvm/lib/Target/DirectX/DXILOpBuilder.cpp llvm/lib/Target/DirectX/DXILOpLowering.cpp

The following files introduce new uses of undef:

• llvm/lib/Target/DirectX/DXILOpLowering.cpp
• llvm/test/CodeGen/DirectX/bufferGetDimensions.ll

Undef is now deprecated and should only be used in the rare cases where no replacement is possible. For example, a load of uninitialized memory yields undef. You should use poison values for placeholders instead.

In tests, avoid using undef and having tests that trigger undefined behavior. If you need an operand with some unimportant value, you can add a new argument to the function and use that instead.

For example, this is considered a bad practice:

define void @fn() {
  ...
  br i1 undef, ...
}

Please use the following instead:

define void @fn(i1 %cond) {
  ...
  br i1 %cond, ...
}

Please refer to the Undefined Behavior Manual for more information.