[clang] [llvm] [HLSL][DirectX] Finish WaveGetLaneCount implementation (PR #156338)

[Nathan Gauër via llvm-commits]

https://github.com/Keenuts created https://github.com/llvm/llvm-project/pull/156338

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

Fixes #99159

>From fd5bfbba2de4ae25336487b9778d56b5292d60ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nathan=20Gau=C3=ABr?= <brioche at google.com>
Date: Mon, 1 Sep 2025 15:50:11 +0200
Subject: [PATCH] [HLSL][DirectX] Finish WaveGetLaneCount implementation

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

Fixes #99159
---
 clang/test/SemaHLSL/WaveGetLaneCount-sm5.hlsl     |  9 +++++++++
 clang/test/SemaHLSL/WaveGetLaneCount-sm6.hlsl     |  9 +++++++++
 llvm/lib/Target/DirectX/DXIL.td                   |  9 +++++++++
 llvm/lib/Target/DirectX/DXILShaderFlags.cpp       |  1 +
 llvm/test/CodeGen/DirectX/ShaderFlags/wave-ops.ll |  7 +++++++
 llvm/test/CodeGen/DirectX/WaveGetLaneCount.ll     | 13 +++++++++++++
 6 files changed, 48 insertions(+)
 create mode 100644 clang/test/SemaHLSL/WaveGetLaneCount-sm5.hlsl
 create mode 100644 clang/test/SemaHLSL/WaveGetLaneCount-sm6.hlsl
 create mode 100644 llvm/test/CodeGen/DirectX/WaveGetLaneCount.ll

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 at 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()
+