[HLSL][DirectX] Finish WaveGetLaneCount implementation

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
+}

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();
+}
+

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";

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:

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]]

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