[HLSL][DirectX] Finish WaveGetLaneCount implementation

[Keenuts]

Some testing was lacking, as well as DirectX lowering of the intrinsic.

Fixes #99159

[llvmbot]

@llvm/pr-subscribers-clang
@llvm/pr-subscribers-hlsl

@llvm/pr-subscribers-backend-directx

Author: Nathan Gauër (Keenuts)

Changes

Some testing was lacking, as well as DirectX lowering of the intrinsic.

Fixes #99159

---

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

6 Files Affected:

• (added) clang/test/SemaHLSL/WaveGetLaneCount-sm5.hlsl (+9)
• (added) clang/test/SemaHLSL/WaveGetLaneCount-sm6.hlsl (+9)
• (modified) llvm/lib/Target/DirectX/DXIL.td (+9)
• (modified) llvm/lib/Target/DirectX/DXILShaderFlags.cpp (+1)
• (modified) llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll (+7)
• (added) llvm/test/CodeGen/DirectX/WaveGetLaneCount.ll (+13)

diff --git a/clang/test/SemaHLSL/WaveGetLaneCount-sm5.hlsl b/clang/test/SemaHLSL/WaveGetLaneCount-sm5.hlsl
new file mode 100644
index 0000000000000..7a12dba62367e
--- /dev/null
+++ b/clang/test/SemaHLSL/WaveGetLaneCount-sm5.hlsl
@@ -0,0 +1,9 @@
+// RUN: %clang_cc1 -finclude-default-header -triple dxil-pc-shadermodel5.1-library -verify %s
+
+[shader("compute")]
+[numthreads(8,8,1)]
+unsigned foo() {
+    // expected-error@#site {{'WaveGetLaneCount' is only available on Shader Model 6.0 or newer}}
+    // expected-note@hlsl/hlsl_alias_intrinsics.h:* {{'WaveGetLaneCount' has been marked as being introduced in Shader Model 6.0 here, but the deployment target is Shader Model 5.1}}
+    return hlsl::WaveGetLaneCount(); // #site
+}
diff --git a/clang/test/SemaHLSL/WaveGetLaneCount-sm6.hlsl b/clang/test/SemaHLSL/WaveGetLaneCount-sm6.hlsl
new file mode 100644
index 0000000000000..b988b1cb4ad54
--- /dev/null
+++ b/clang/test/SemaHLSL/WaveGetLaneCount-sm6.hlsl
@@ -0,0 +1,9 @@
+// RUN: %clang_cc1 -finclude-default-header -triple dxil-pc-shadermodel6.0-library -verify %s
+
+// expected-no-diagnostics
+[shader("compute")]
+[numthreads(8,8,1)]
+unsigned foo() {
+    return hlsl::WaveGetLaneCount();
+}
+
diff --git a/llvm/lib/Target/DirectX/DXIL.td b/llvm/lib/Target/DirectX/DXIL.td
index c65ead45e2c7e..c9977e3bfdbf4 100644
--- a/llvm/lib/Target/DirectX/DXIL.td
+++ b/llvm/lib/Target/DirectX/DXIL.td
@@ -998,6 +998,15 @@ def WaveGetLaneIndex : DXILOp<111, waveGetLaneIndex> {
   let attributes = [Attributes<DXIL1_0, [ReadOnly]>];
 }
 
+def WaveGetLaneCount : DXILOp<112, waveGetLaneCount> {
+  let Doc = "returns the number of lanes in a wave on this architecture";
+  let intrinsics = [IntrinSelect<int_dx_wave_get_lane_count>];
+  let arguments = [];
+  let result = Int32Ty;
+  let stages = [Stages<DXIL1_0, [all_stages]>];
+  let attributes = [Attributes<DXIL1_0, [ReadOnly]>];
+}
+
 def WaveActiveAnyTrue : DXILOp<113, waveAnyTrue> {
   let Doc = "returns true if the expression is true in any of the active lanes "
             "in the current wave";
diff --git a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
index e7e7f2ce66ae8..58a94f7b1cb9d 100644
--- a/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
+++ b/llvm/lib/Target/DirectX/DXILShaderFlags.cpp
@@ -85,6 +85,7 @@ static bool checkWaveOps(Intrinsic::ID IID) {
     return false;
   case Intrinsic::dx_wave_is_first_lane:
   case Intrinsic::dx_wave_getlaneindex:
+  case Intrinsic::dx_wave_get_lane_count:
   case Intrinsic::dx_wave_any:
   case Intrinsic::dx_wave_all:
   case Intrinsic::dx_wave_readlane:
diff --git a/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll b/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll
index 7a876f67615cd..fab7657f862d8 100644
--- a/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll
+++ b/llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll
@@ -27,6 +27,13 @@ entry:
   ret i32 %ret
 }
 
+define noundef i32 @wave_get_lane_count() {
+entry:
+  ; CHECK: Function wave_get_lane_count : [[WAVE_FLAG]]
+  %ret = call i32 @llvm.dx.wave.get.lane.count()
+  ret i32 %ret
+}
+
 define noundef i1 @wave_any(i1 %x) {
 entry:
   ; CHECK: Function wave_any : [[WAVE_FLAG]]
diff --git a/llvm/test/CodeGen/DirectX/WaveGetLaneCount.ll b/llvm/test/CodeGen/DirectX/WaveGetLaneCount.ll
new file mode 100644
index 0000000000000..784d5781f755e
--- /dev/null
+++ b/llvm/test/CodeGen/DirectX/WaveGetLaneCount.ll
@@ -0,0 +1,13 @@
+; RUN: opt -S  -dxil-op-lower  -mtriple=dxil-pc-shadermodel6.3-compute %s | FileCheck %s
+
+define void @main() {
+entry:
+; CHECK: call i32 @dx.op.waveGetLaneCount(i32 112) #[[#ATTR:]]
+  %0 = call i32 @llvm.dx.wave.get.lane.count()
+  ret void
+}
+
+; CHECK: attributes #[[#ATTR]] = {{{.*}} memory(read) {{.*}}}
+
+declare i32 @llvm.dx.wave.get.lane.count()
+

[farzonl]

see line 66 seems like you need to remove that commented out case statement since you are adding it here.

[farzonl]

lgtm