[clang][HLSL] Add WaveIsFirstLane() intrinsic (#103299)

This commits add the WaveIsFirstLane() hlsl intrinsinc. This intrinsic
uses the convergence intrinsincs for the SPIR-V backend. On the DXIL
side, I'm not sure what the strategy is for convergence, so I
implemented that like in DXC: a normal builtin function.

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

clang/include/clang/Basic/Builtins.td

@@ -4679,6 +4679,12 @@ def HLSLWaveGetLaneIndex : LangBuiltin<"HLSL_LANG"> {
   let Prototype = "unsigned int()";
 }
 
+def HLSLWaveIsFirstLane : LangBuiltin<"HLSL_LANG"> {
+  let Spellings = ["__builtin_hlsl_wave_is_first_lane"];
+  let Attributes = [NoThrow, Const];
+  let Prototype = "bool()";
+}
+
 def HLSLClamp : LangBuiltin<"HLSL_LANG"> {
   let Spellings = ["__builtin_hlsl_elementwise_clamp"];
   let Attributes = [NoThrow, Const];

clang/lib/CodeGen/CGBuiltin.cpp

@@ -18723,6 +18723,10 @@ case Builtin::BI__builtin_hlsl_elementwise_isinf: {
         llvm::FunctionType::get(IntTy, {}, false), "__hlsl_wave_get_lane_index",
         {}, false, true));
   }
+  case Builtin::BI__builtin_hlsl_wave_is_first_lane: {
+    Intrinsic::ID ID = CGM.getHLSLRuntime().getWaveIsFirstLaneIntrinsic();
+    return EmitRuntimeCall(Intrinsic::getDeclaration(&CGM.getModule(), ID));
+  }
   }
   return nullptr;
 }

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -84,6 +84,7 @@ class CGHLSLRuntime {
   GENERATE_HLSL_INTRINSIC_FUNCTION(FDot, fdot)
   GENERATE_HLSL_INTRINSIC_FUNCTION(SDot, sdot)
   GENERATE_HLSL_INTRINSIC_FUNCTION(UDot, udot)
+  GENERATE_HLSL_INTRINSIC_FUNCTION(WaveIsFirstLane, wave_is_first_lane)
 
   //===----------------------------------------------------------------------===//
   // End of reserved area for HLSL intrinsic getters.

clang/lib/Headers/hlsl/hlsl_intrinsics.h

@@ -1796,5 +1796,9 @@ _HLSL_AVAILABILITY(shadermodel, 6.0)
 _HLSL_BUILTIN_ALIAS(__builtin_hlsl_wave_get_lane_index)
 __attribute__((convergent)) uint WaveGetLaneIndex();
 
+_HLSL_AVAILABILITY(shadermodel, 6.0)
+_HLSL_BUILTIN_ALIAS(__builtin_hlsl_wave_is_first_lane)
+__attribute__((convergent)) bool WaveIsFirstLane();
+
 } // namespace hlsl
 #endif //_HLSL_HLSL_INTRINSICS_H_

clang/test/CodeGenHLSL/builtins/wave_is_first_lane.hlsl

@@ -0,0 +1,34 @@
+// RUN: %clang_cc1 -std=hlsl2021 -finclude-default-header -x hlsl -triple   \
+// RUN:   spirv-pc-vulkan-compute %s -emit-llvm -disable-llvm-passes -o - | \
+// RUN:   FileCheck %s --check-prefixes=CHECK,CHECK-SPIRV
+// RUN: %clang_cc1 -std=hlsl2021 -finclude-default-header -x hlsl -triple   \
+// RUN:   dxil-pc-shadermodel6.3-compute %s -emit-llvm -disable-llvm-passes -o - | \
+// RUN:   FileCheck %s --check-prefixes=CHECK,CHECK-DXIL
+
+[numthreads(1, 1, 1)]
+void main() {
+// CHECK-SPIRV: %[[#entry_tok:]] = call token @llvm.experimental.convergence.entry()
+
+// CHECK-SPIRV: %[[#loop_tok:]] = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %[[#entry_tok]]) ]
+  while (true) {
+
+// CHECK-DXIL:  %[[#]] = call i1 @llvm.dx.wave.is.first.lane()
+// CHECK-SPIRV: %[[#]] = call i1 @llvm.spv.wave.is.first.lane()
+// CHECK-SPIRV-SAME: [ "convergencectrl"(token %[[#loop_tok]]) ]
+    if (WaveIsFirstLane()) {
+      break;
+    }
+  }
+
+// CHECK-DXIL:  %[[#]] = call i1 @llvm.dx.wave.is.first.lane()
+// CHECK-SPIRV: %[[#]] = call i1 @llvm.spv.wave.is.first.lane()
+// CHECK-SPIRV-SAME: [ "convergencectrl"(token %[[#entry_tok]]) ]
+  if (WaveIsFirstLane()) {
+    return;
+  }
+}
+
+// CHECK-DXIL:  i1 @llvm.dx.wave.is.first.lane() #[[#attr:]]
+// CHECK-SPIRV: i1 @llvm.spv.wave.is.first.lane() #[[#attr:]]
+
+// CHECK: attributes #[[#attr]] = {{{.*}} convergent {{.*}}}

llvm/include/llvm/IR/IntrinsicsDirectX.td

@@ -79,4 +79,6 @@ def int_dx_umad : DefaultAttrsIntrinsic<[llvm_anyint_ty], [LLVMMatchType<0>, LLV
 def int_dx_normalize : DefaultAttrsIntrinsic<[LLVMMatchType<0>], [llvm_anyfloat_ty]>;
 def int_dx_rcp  : DefaultAttrsIntrinsic<[llvm_anyfloat_ty], [LLVMMatchType<0>]>;
 def int_dx_rsqrt  : DefaultAttrsIntrinsic<[llvm_anyfloat_ty], [LLVMMatchType<0>]>;
+
+def int_dx_wave_is_first_lane : DefaultAttrsIntrinsic<[llvm_i1_ty], [], [IntrConvergent]>;
 }

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -79,4 +79,5 @@ let TargetPrefix = "spv" in {
     DefaultAttrsIntrinsic<[LLVMVectorElementType<0>],
     [llvm_anyint_ty, LLVMScalarOrSameVectorWidth<0, LLVMVectorElementType<0>>],
     [IntrNoMem, Commutative] >;
+  def int_spv_wave_is_first_lane : DefaultAttrsIntrinsic<[llvm_i1_ty], [], [IntrConvergent]>;
 }

llvm/lib/Target/DirectX/DXIL.td

@@ -746,3 +746,12 @@ def CreateHandleFromBinding : DXILOp<218, createHandleFromBinding> {
   let result = HandleTy;
   let stages = [Stages<DXIL1_6, [all_stages]>];
 }
+
+def WaveIsFirstLane :  DXILOp<110, waveIsFirstLane> {
+  let Doc = "returns 1 for the first lane in the wave";
+  let LLVMIntrinsic = int_dx_wave_is_first_lane;
+  let arguments = [];
+  let result = Int1Ty;
+  let stages = [Stages<DXIL1_0, [all_stages]>];
+  let attributes = [Attributes<DXIL1_0, [ReadNone]>];
+}

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -2351,6 +2351,14 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
   } break;
   case Intrinsic::spv_saturate:
     return selectSaturate(ResVReg, ResType, I);
+  case Intrinsic::spv_wave_is_first_lane: {
+    SPIRVType *IntTy = GR.getOrCreateSPIRVIntegerType(32, I, TII);
+    return BuildMI(BB, I, I.getDebugLoc(),
+                   TII.get(SPIRV::OpGroupNonUniformElect))
+        .addDef(ResVReg)
+        .addUse(GR.getSPIRVTypeID(ResType))
+        .addUse(GR.getOrCreateConstInt(3, I, IntTy, TII));
+  }
   default: {
     std::string DiagMsg;
     raw_string_ostream OS(DiagMsg);

llvm/lib/Target/SPIRV/SPIRVStripConvergentIntrinsics.cpp

@@ -41,31 +41,40 @@ class SPIRVStripConvergentIntrinsics : public FunctionPass {
   virtual bool runOnFunction(Function &F) override {
     DenseSet<Instruction *> ToRemove;
 
+    // Is the instruction is a convergent intrinsic, add it to kill-list and
+    // returns true. Returns false otherwise.
+    auto CleanupIntrinsic = [&](IntrinsicInst *II) {
+      if (II->getIntrinsicID() != Intrinsic::experimental_convergence_entry &&
+          II->getIntrinsicID() != Intrinsic::experimental_convergence_loop &&
+          II->getIntrinsicID() != Intrinsic::experimental_convergence_anchor)
+        return false;
+
+      II->replaceAllUsesWith(UndefValue::get(II->getType()));
+      ToRemove.insert(II);
+      return true;
+    };
+
+    // Replace the given CallInst by a similar CallInst with no convergencectrl
+    // attribute.
+    auto CleanupCall = [&](CallInst *CI) {
+      auto OB = CI->getOperandBundle(LLVMContext::OB_convergencectrl);
+      if (!OB.has_value())
+        return;
+
+      auto *NewCall = CallBase::removeOperandBundle(
+          CI, LLVMContext::OB_convergencectrl, CI);
+      NewCall->copyMetadata(*CI);
+      CI->replaceAllUsesWith(NewCall);
+      ToRemove.insert(CI);
+    };
+
     for (BasicBlock &BB : F) {
       for (Instruction &I : BB) {
-        if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
-          if (II->getIntrinsicID() !=
-                  Intrinsic::experimental_convergence_entry &&
-              II->getIntrinsicID() !=
-                  Intrinsic::experimental_convergence_loop &&
-              II->getIntrinsicID() !=
-                  Intrinsic::experimental_convergence_anchor) {
+        if (auto *II = dyn_cast<IntrinsicInst>(&I))
+          if (CleanupIntrinsic(II))
             continue;
-          }
-
-          II->replaceAllUsesWith(UndefValue::get(II->getType()));
-          ToRemove.insert(II);
-        } else if (auto *CI = dyn_cast<CallInst>(&I)) {
-          auto OB = CI->getOperandBundle(LLVMContext::OB_convergencectrl);
-          if (!OB.has_value())
-            continue;
-
-          auto *NewCall = CallBase::removeOperandBundle(
-              CI, LLVMContext::OB_convergencectrl, CI);
-          NewCall->copyMetadata(*CI);
-          CI->replaceAllUsesWith(NewCall);
-          ToRemove.insert(CI);
-        }
+        if (auto *CI = dyn_cast<CallInst>(&I))
+          CleanupCall(CI);
       }
     }