[llvm] 5742dc4 - [DirectX][ShaderFlags] Add analysis for `WaveOps` flag (#118140)

[via llvm-commits]

Author: Finn Plummer
Date: 2025-02-14T10:40:47-08:00
New Revision: 5742dc44e2fb571952990b4226b10476e15ac479

URL: https://github.com/llvm/llvm-project/commit/5742dc44e2fb571952990b4226b10476e15ac479
DIFF: https://github.com/llvm/llvm-project/commit/5742dc44e2fb571952990b4226b10476e15ac479.diff

LOG: [DirectX][ShaderFlags] Add analysis for `WaveOps` flag (#118140)

- Check each call instruction for a `WaveOp` intrinsic and set the
`WaveOps` flag if this is true for any intrinsic, Done in
DXILShaderFlags.cpp

Resolves #114565

Added: 
    llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll

Modified: 
    llvm/lib/Target/DirectX/DXILShaderFlags.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
index 453aacf49abb3..e1cf516758663 100644
--- a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
+++ b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
@@ -32,6 +32,43 @@
 using namespace llvm;
 using namespace llvm::dxil;
 
+static bool checkWaveOps(Intrinsic::ID IID) {
+  // Currently unsupported intrinsics
+  // case Intrinsic::dx_wave_getlanecount:
+  // case Intrinsic::dx_wave_allequal:
+  // case Intrinsic::dx_wave_ballot:
+  // case Intrinsic::dx_wave_readfirst:
+  // case Intrinsic::dx_wave_reduce.and:
+  // case Intrinsic::dx_wave_reduce.or:
+  // case Intrinsic::dx_wave_reduce.xor:
+  // case Intrinsic::dx_wave_prefixop:
+  // case Intrinsic::dx_quad.readat:
+  // case Intrinsic::dx_quad.readacrossx:
+  // case Intrinsic::dx_quad.readacrossy:
+  // case Intrinsic::dx_quad.readacrossdiagonal:
+  // case Intrinsic::dx_wave_prefixballot:
+  // case Intrinsic::dx_wave_match:
+  // case Intrinsic::dx_wavemulti.*:
+  // case Intrinsic::dx_wavemulti.ballot:
+  // case Intrinsic::dx_quad.vote:
+  switch (IID) {
+  default:
+    return false;
+  case Intrinsic::dx_wave_is_first_lane:
+  case Intrinsic::dx_wave_getlaneindex:
+  case Intrinsic::dx_wave_any:
+  case Intrinsic::dx_wave_all:
+  case Intrinsic::dx_wave_readlane:
+  case Intrinsic::dx_wave_active_countbits:
+  // Wave Active Op Variants
+  case Intrinsic::dx_wave_reduce_sum:
+  case Intrinsic::dx_wave_reduce_usum:
+  case Intrinsic::dx_wave_reduce_max:
+  case Intrinsic::dx_wave_reduce_umax:
+    return true;
+  }
+}
+
 /// Update the shader flags mask based on the given instruction.
 /// \param CSF Shader flags mask to update.
 /// \param I Instruction to check.
@@ -94,6 +131,8 @@ void ModuleShaderFlags::updateFunctionFlags(ComputedShaderFlags &CSF,
 
     // TODO: Set DX11_1_DoubleExtensions if I is a call to DXIL intrinsic
     // DXIL::Opcode::Fma https://github.com/llvm/llvm-project/issues/114554
+
+    CSF.WaveOps |= checkWaveOps(CI->getIntrinsicID());
   }
 }
 

diff  --git a/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll
new file mode 100644
index 0000000000000..7a876f67615cd
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll
@@ -0,0 +1,84 @@
+; RUN: opt -S --passes="print-dx-shader-flags" 2>&1 %s | FileCheck %s
+;
+; Test that we have the correct shader flags to indicate that there are wave
+; ops set at the module level
+;
+; CHECK: ; Shader Flags Value: [[WAVE_FLAG:0x00080000]]
+; CHECK: ; Note: shader requires additional functionality:
+; CHECK-NEXT: ;       Wave level operations
+; CHECK-NEXT: ; Note: extra DXIL module flags:
+
+target triple = "dxil-pc-shadermodel6.7-library"
+
+; Test the indiviual ops that they have the same Shader Wave flag at the
+; function level to ensure that each op is setting it accordingly
+
+define noundef i1 @wave_is_first_lane() {
+entry:
+  ; CHECK: Function wave_is_first_lane : [[WAVE_FLAG]]
+  %ret = call i1 @llvm.dx.wave.is.first.lane()
+  ret i1 %ret
+}
+
+define noundef i32 @wave_getlaneindex() {
+entry:
+  ; CHECK: Function wave_getlaneindex : [[WAVE_FLAG]]
+  %ret = call i32 @llvm.dx.wave.getlaneindex()
+  ret i32 %ret
+}
+
+define noundef i1 @wave_any(i1 %x) {
+entry:
+  ; CHECK: Function wave_any : [[WAVE_FLAG]]
+  %ret = call i1 @llvm.dx.wave.any(i1 %x)
+  ret i1 %ret
+}
+
+define noundef i1 @wave_all(i1 %x) {
+entry:
+  ; CHECK: Function wave_all : [[WAVE_FLAG]]
+  %ret = call i1 @llvm.dx.wave.all(i1 %x)
+  ret i1 %ret
+}
+
+define noundef i1 @wave_readlane(i1 %x, i32 %idx) {
+entry:
+  ; CHECK: Function wave_readlane : [[WAVE_FLAG]]
+  %ret = call i1 @llvm.dx.wave.readlane.i1(i1 %x, i32 %idx)
+  ret i1 %ret
+}
+
+define noundef i32 @wave_reduce_sum(i32 noundef %x) {
+entry:
+  ; CHECK: Function wave_reduce_sum : [[WAVE_FLAG]]
+  %ret = call i32 @llvm.dx.wave.reduce.sum.i32(i32 %x)
+  ret i32 %ret
+}
+
+define noundef i32 @wave_reduce_usum(i32 noundef %x) {
+entry:
+  ; CHECK: Function wave_reduce_usum : [[WAVE_FLAG]]
+  %ret = call i32 @llvm.dx.wave.reduce.usum.i32(i32 %x)
+  ret i32 %ret
+}
+
+define noundef i32 @wave_reduce_max(i32 noundef %x) {
+entry:
+  ; CHECK: Function wave_reduce_max : [[WAVE_FLAG]]
+  %ret = call i32 @llvm.dx.wave.reduce.max.i32(i32 %x)
+  ret i32 %ret
+}
+
+define noundef i32 @wave_reduce_umax(i32 noundef %x) {
+entry:
+  ; CHECK: Function wave_reduce_umax : [[WAVE_FLAG]]
+  %ret = call i32 @llvm.dx.wave.reduce.umax.i32(i32 %x)
+  ret i32 %ret
+}
+
+define void @wave_active_countbits(i1 %expr) {
+entry:
+  ; CHECK: Function wave_active_countbits : [[WAVE_FLAG]]
+  %0 = call i32 @llvm.dx.wave.active.countbits(i1 %expr)
+  ret void
+}