[llvm] [DirectX] Handle <1 x ...> loads in DXILResourceAccess (PR #137076)

[Justin Bogner via llvm-commits]

https://github.com/bogner updated https://github.com/llvm/llvm-project/pull/137076

>From d90e8202e623c0f164e5b0a3ed53b053bc8a8fb2 Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail at justinbogner.com>
Date: Wed, 23 Apr 2025 15:46:46 -0600
Subject: [PATCH 1/2] [DirectX] Handle <1 x ...> loads in DXILResourceAccess

We can end up with loads of single element vectors when we have scalar
values, because the vectorizer may introduce these to use ops like
shufflevector in some cases. Make sure we're maintaining the correct
type when translating these into resource load operations.

Fixes #136409.
---
 .../lib/Target/DirectX/DXILResourceAccess.cpp |  7 +++
 .../ResourceAccess/load_typedbuffer.ll        | 48 ++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/DirectX/DXILResourceAccess.cpp b/llvm/lib/Target/DirectX/DXILResourceAccess.cpp
index 3b8f7140d3122..16337f1237e00 100644
--- a/llvm/lib/Target/DirectX/DXILResourceAccess.cpp
+++ b/llvm/lib/Target/DirectX/DXILResourceAccess.cpp
@@ -143,6 +143,13 @@ static void createTypedBufferLoad(IntrinsicInst *II, LoadInst *LI,
   if (Offset)
     V = Builder.CreateExtractElement(V, Offset);
 
+  // If we loaded a <1 x ...> instead of a scalar (presumably to feed a
+  // shufflevector), then make sure we're maintaining the resulting type.
+  if (auto *VT = dyn_cast<FixedVectorType>(LI->getType()))
+    if (VT->getNumElements() == 1 && !isa<FixedVectorType>(V->getType()))
+      V = Builder.CreateInsertElement(PoisonValue::get(VT), V,
+                                      Builder.getInt32(0));
+
   LI->replaceAllUsesWith(V);
 }
 
diff --git a/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll b/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll
index 8769e6ec66d8e..7d8099ef0d871 100644
--- a/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll
+++ b/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll
@@ -4,8 +4,9 @@ target triple = "dxil-pc-shadermodel6.6-compute"
 
 declare void @use_float4(<4 x float>)
 declare void @use_float(float)
+declare void @use_float1(<1 x float>)
 
-; CHECK-LABEL: define void @load_float4
+; CHECK-LABEL: define void @load_float4(
 define void @load_float4(i32 %index, i32 %elemindex) {
   %buffer = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
       @llvm.dx.resource.handlefrombinding.tdx.TypedBuffer_v4f32_1_0_0(
@@ -35,3 +36,48 @@ define void @load_float4(i32 %index, i32 %elemindex) {
 
   ret void
 }
+
+; CHECK-LABEL: define void @load_float(
+define void @load_float(i32 %index, i32 %elemindex) {
+  %buffer = call target("dx.TypedBuffer", float, 1, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+
+  ; CHECK-NOT: @llvm.dx.resource.getpointer
+  %ptr = call ptr @llvm.dx.resource.getpointer(
+      target("dx.TypedBuffer", float, 1, 0, 0) %buffer, i32 %index)
+
+  ; CHECK: %[[LOAD:.*]] = call { float, i1 } @llvm.dx.resource.load.typedbuffer.f32.tdx.TypedBuffer_f32_1_0_0t(target("dx.TypedBuffer", float, 1, 0, 0) %buffer, i32 %index)
+  ; CHECK: %[[VALUE:.*]] = extractvalue { float, i1 } %[[LOAD]], 0
+  ; CHECK: call void @use_float(float %[[VALUE]])
+  %data = load float, ptr %ptr
+  call void @use_float(float %data)
+
+  ; CHECK: %[[LOAD:.*]] = call { float, i1 } @llvm.dx.resource.load.typedbuffer.f32.tdx.TypedBuffer_f32_1_0_0t(target("dx.TypedBuffer", float, 1, 0, 0) %buffer, i32 %index)
+  ; CHECK: %[[VALUE:.*]] = extractvalue { float, i1 } %[[LOAD]], 0
+  ; CHECK: %[[VEC:.*]] = insertelement <1 x float> poison, float %[[VALUE]], i32 0
+  ; CHECK: call void @use_float1(<1 x float> %[[VEC]])
+  %vec_data = load <1 x float>, ptr %ptr
+  call void @use_float1(<1 x float> %vec_data)
+
+  ret void
+}
+
+; CHECK-LABEL: define void @load_float1(
+define void @load_float1(i32 %index) {
+  %buffer = call target("dx.TypedBuffer", <1 x float>, 1, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+
+  ; CHECK-NOT: @llvm.dx.resource.getpointer
+  %ptr = call ptr @llvm.dx.resource.getpointer(
+      target("dx.TypedBuffer", <1 x float>, 1, 0, 0) %buffer, i32 %index)
+
+  ; CHECK: %[[LOAD:.*]] = call { <1 x float>, i1 } @llvm.dx.resource.load.typedbuffer.v1f32.tdx.TypedBuffer_v1f32_1_0_0t(target("dx.TypedBuffer", <1 x float>, 1, 0, 0) %buffer, i32 %index)
+  ; CHECK: %[[VALUE:.*]] = extractvalue { <1 x float>, i1 } %[[LOAD]], 0
+  ; CHECK: %[[SHUF:.*]] = shufflevector <1 x float> %[[VALUE]], <1 x float> poison, <4 x i32> zeroinitializer
+  ; CHECK: call void @use_float4(<4 x float> %[[SHUF]])
+  %load = load <1 x float>, ptr %ptr
+  %shuf = shufflevector <1 x float> %load, <1 x float> poison, <4 x i32> zeroinitializer
+  call void @use_float4(<4 x float> %shuf)
+
+  ret void
+}

>From 2e2c91812d8395e0b27b9b7eef2486345c2780b6 Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail at justinbogner.com>
Date: Thu, 24 Apr 2025 14:16:16 -0600
Subject: [PATCH 2/2] fixup! [DirectX] Handle <1 x ...> loads in
 DXILResourceAccess

---
 llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll b/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll
index 7d8099ef0d871..b51b9ae4b4a11 100644
--- a/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll
+++ b/llvm/test/CodeGen/DirectX/ResourceAccess/load_typedbuffer.ll
@@ -38,7 +38,7 @@ define void @load_float4(i32 %index, i32 %elemindex) {
 }
 
 ; CHECK-LABEL: define void @load_float(
-define void @load_float(i32 %index, i32 %elemindex) {
+define void @load_float(i32 %index) {
   %buffer = call target("dx.TypedBuffer", float, 1, 0, 0)
       @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)