[llvm] [DirectX] Set the EnableRawAndStructuredBuffers shader flag (PR #122667)

[Justin Bogner via llvm-commits]

https://github.com/bogner created https://github.com/llvm/llvm-project/pull/122667

When raw or structured buffers are used, we need to set the DXIL flag saying so.

Fixes #122663.

>From 328d1d311af1da93d6465fd1f42cb561ee0dc8fa Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail at justinbogner.com>
Date: Fri, 20 Dec 2024 17:44:00 -0700
Subject: [PATCH] [DirectX] Set the EnableRawAndStructuredBuffers shader flag

When raw or structured buffers are used, we need to set the DXIL flag
saying so.

Fixes #122663.
---
 .../BinaryFormat/DXContainerConstants.def     |  2 +-
 llvm/lib/Target/DirectX/DXILShaderFlags.cpp   | 11 +++++
 .../ShaderFlags/raw-and-structured-buffers.ll | 42 +++++++++++++++++++
 3 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/DirectX/ShaderFlags/raw-and-structured-buffers.ll

diff --git a/llvm/include/llvm/BinaryFormat/DXContainerConstants.def b/llvm/include/llvm/BinaryFormat/DXContainerConstants.def
index 1aacbb2f65b27f..96d4499c9cadc9 100644
--- a/llvm/include/llvm/BinaryFormat/DXContainerConstants.def
+++ b/llvm/include/llvm/BinaryFormat/DXContainerConstants.def
@@ -58,7 +58,7 @@ SHADER_FEATURE_FLAG(31, 36, NextUnusedBit, "Next reserved shader flag bit (not a
 DXIL_MODULE_FLAG( 0,  DisableOptimizations,   "D3D11_1_SB_GLOBAL_FLAG_SKIP_OPTIMIZATION")
 DXIL_MODULE_FLAG( 1,  DisableMathRefactoring, "D3D10_SB_GLOBAL_FLAG_REFACTORING_ALLOWED")
 DXIL_MODULE_FLAG( 3,  ForceEarlyDepthStencil, "D3D11_SB_GLOBAL_FLAG_FORCE_EARLY_DEPTH_STENCIL")
-DXIL_MODULE_FLAG( 4,  EnableRawAndStructuredBuffers, "D3D11_SB_GLOBAL_FLAG_ENABLE_RAW_AND_STRUCTURED_BUFFERS")
+DXIL_MODULE_FLAG( 4,  EnableRawAndStructuredBuffers, "Raw and Structured buffers")
 DXIL_MODULE_FLAG( 5,  LowPrecisionPresent, "D3D11_1_SB_GLOBAL_FLAG_ENABLE_MINIMUM_PRECISION")
 DXIL_MODULE_FLAG( 8,  AllResourcesBound, "D3D12_SB_GLOBAL_FLAG_ALL_RESOURCES_BOUND")
 DXIL_MODULE_FLAG(23,  UseNativeLowPrecision, "Native 16bit types enabled")
diff --git a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
index 2edfc707ce6c79..36e9119363f5c7 100644
--- a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
+++ b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
@@ -54,11 +54,22 @@ static void updateFunctionFlags(ComputedShaderFlags &CSF, const Instruction &I,
     switch (II->getIntrinsicID()) {
     default:
       break;
+    case Intrinsic::dx_resource_handlefrombinding:
+      switch (DRTM[cast<TargetExtType>(II->getType())].getResourceKind()) {
+      case dxil::ResourceKind::StructuredBuffer:
+      case dxil::ResourceKind::RawBuffer:
+        CSF.EnableRawAndStructuredBuffers = true;
+        break;
+      default:
+        break;
+      }
+      break;
     case Intrinsic::dx_resource_load_typedbuffer: {
       dxil::ResourceTypeInfo &RTI =
           DRTM[cast<TargetExtType>(II->getArgOperand(0)->getType())];
       if (RTI.isTyped())
         CSF.TypedUAVLoadAdditionalFormats |= RTI.getTyped().ElementCount > 1;
+      break;
     }
     }
   }
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/raw-and-structured-buffers.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/raw-and-structured-buffers.ll
new file mode 100644
index 00000000000000..9cb83c4e992bad
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/raw-and-structured-buffers.ll
@@ -0,0 +1,42 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+
+; Note: there is no feature flag here (only a module flag), so we don't have an
+; object test.
+
+target triple = "dxil-pc-shadermodel6.7-library"
+
+; CHECK:      Combined Shader Flags for Module
+; CHECK-NEXT: Shader Flags Value: 0x00000010
+
+; CHECK: Note: shader requires additional functionality:
+; CHECK:       Raw and Structured buffers
+
+; CHECK: Function rawbuf : 0x00000010
+define float @rawbuf() "hlsl.export" {
+  %buffer = call target("dx.RawBuffer", i8, 0, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  %load = call {float, i1} @llvm.dx.resource.load.rawbuffer.f32(
+      target("dx.RawBuffer", i8, 0, 0, 0) %buffer, i32 0, i32 0)
+  %data = extractvalue {float, i1} %load, 0
+  ret float %data
+}
+
+; CHECK: Function structbuf : 0x00000010
+define float @structbuf() "hlsl.export" {
+  %buffer = call target("dx.RawBuffer", float, 0, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  %load = call {float, i1} @llvm.dx.resource.load.rawbuffer.f32(
+      target("dx.RawBuffer", float, 0, 0, 0) %buffer, i32 0, i32 0)
+  %data = extractvalue {float, i1} %load, 0
+  ret float %data
+}
+
+; CHECK: Function typedbuf : 0x00000000
+define float @typedbuf(<4 x float> %val) "hlsl.export" {
+  %buffer = call target("dx.TypedBuffer", float, 0, 0, 0)
+      @llvm.dx.resource.handlefrombinding(i32 0, i32 0, i32 1, i32 0, i1 false)
+  %load = call {float, i1} @llvm.dx.resource.load.typedbuffer(
+      target("dx.TypedBuffer", float, 0, 0, 0) %buffer, i32 0)
+  %data = extractvalue {float, i1} %load, 0
+  ret float %data
+}