Add cl_khr_kernel_clock support (#2582)

svenvh · sys-ce-bb · commit a050277734ee · 2024-05-30T06:45:38.000-07:00
Add support for mapping the `cl_khr_kernel_clock` extension builtins to and from the `SPV_KHR_shader_clock` extension. Original commit: KhronosGroup/SPIRV-LLVM-Translator@abf48906dae1ed4
diff --git a/llvm-spirv/lib/SPIRV/OCLToSPIRV.cpp b/llvm-spirv/lib/SPIRV/OCLToSPIRV.cpp
@@ -340,6 +340,10 @@ void OCLToSPIRVBase::visitCallInst(CallInst &CI) {
     visitCallDot(&CI, MangledName, DemangledName);
     return;
   }
+  if (DemangledName.starts_with(kOCLBuiltinName::ClockReadPrefix)) {
+    visitCallClockRead(&CI, MangledName, DemangledName);
+    return;
+  }
   if (DemangledName == kOCLBuiltinName::FMin ||
       DemangledName == kOCLBuiltinName::FMax ||
       DemangledName == kOCLBuiltinName::Min ||
@@ -1324,6 +1328,23 @@ void OCLToSPIRVBase::visitCallDot(CallInst *CI, StringRef MangledName,
   }
 }
 
+void OCLToSPIRVBase::visitCallClockRead(CallInst *CI, StringRef MangledName,
+                                        StringRef DemangledName) {
+  // The builtin returns i64 or <2 x i32>, but both variants are mapped to the
+  // same instruction; hence include the return type.
+  std::string OpName = getSPIRVFuncName(OpReadClockKHR, CI->getType());
+
+  // Scope is part of the OpenCL builtin name.
+  Scope ScopeArg = StringSwitch<Scope>(DemangledName)
+                       .EndsWith("device", ScopeDevice)
+                       .EndsWith("work_group", ScopeWorkgroup)
+                       .EndsWith("sub_group", ScopeSubgroup)
+                       .Default(ScopeMax);
+
+  auto Mutator = mutateCallInst(CI, OpName);
+  Mutator.appendArg(getInt32(M, ScopeArg));
+}
+
 void OCLToSPIRVBase::visitCallScalToVec(CallInst *CI, StringRef MangledName,
                                         StringRef DemangledName) {
   // Check if all arguments have the same type - it's simple case.
diff --git a/llvm-spirv/lib/SPIRV/OCLToSPIRV.h b/llvm-spirv/lib/SPIRV/OCLToSPIRV.h
@@ -217,6 +217,10 @@ class OCLToSPIRVBase : public InstVisitor<OCLToSPIRVBase>, BuiltinCallHelper {
   void visitCallDot(CallInst *CI, StringRef MangledName,
                     StringRef DemangledName);
 
+  /// Transform clock_read_* calls to OpReadClockKHR instructions.
+  void visitCallClockRead(CallInst *CI, StringRef MangledName,
+                          StringRef DemangledName);
+
   /// Fixes for built-in functions with vector+scalar arguments that are
   /// translated to the SPIR-V instructions where all arguments must have the
   /// same type.
diff --git a/llvm-spirv/lib/SPIRV/OCLUtil.h b/llvm-spirv/lib/SPIRV/OCLUtil.h
@@ -237,6 +237,7 @@ const static char AtomicInit[] = "atomic_init";
 const static char AtomicWorkItemFence[] = "atomic_work_item_fence";
 const static char Barrier[] = "barrier";
 const static char Clamp[] = "clamp";
+const static char ClockReadPrefix[] = "clock_read_";
 const static char ConvertPrefix[] = "convert_";
 const static char Dot[] = "dot";
 const static char DotAccSat[] = "dot_acc_sat";
diff --git a/llvm-spirv/lib/SPIRV/SPIRVToOCL.cpp b/llvm-spirv/lib/SPIRV/SPIRVToOCL.cpp
@@ -210,6 +210,10 @@ void SPIRVToOCLBase::visitCallInst(CallInst &CI) {
       visitCallSPIRVRelational(&CI, OC);
     return;
   }
+  if (OC == OpReadClockKHR) {
+    visitCallSPIRVReadClockKHR(&CI);
+    return;
+  }
   if (OC == internal::OpConvertFToBF16INTEL ||
       OC == internal::OpConvertBF16ToFINTEL) {
     visitCallSPIRVBFloat16Conversions(&CI, OC);
@@ -1021,6 +1025,33 @@ void SPIRVToOCLBase::visitCallSPIRVRelational(CallInst *CI, Op OC) {
       });
 }
 
+void SPIRVToOCLBase::visitCallSPIRVReadClockKHR(CallInst *CI) {
+  std::ostringstream Name;
+  Name << "clock_read_";
+
+  if (CI->getType()->isVectorTy())
+    Name << "hilo_";
+
+  // Encode the scope (taken from the argument) in the function name.
+  ConstantInt *ScopeOp = cast<ConstantInt>(CI->getArgOperand(0));
+  switch (static_cast<Scope>(ScopeOp->getZExtValue())) {
+  case ScopeDevice:
+    Name << "device";
+    break;
+  case ScopeWorkgroup:
+    Name << "work_group";
+    break;
+  case ScopeSubgroup:
+    Name << "sub_group";
+    break;
+  default:
+    break;
+  }
+
+  auto Mutator = mutateCallInst(CI, Name.str());
+  Mutator.removeArg(0);
+}
+
 std::string SPIRVToOCLBase::getGroupBuiltinPrefix(CallInst *CI) {
   std::string Prefix;
   auto ES = getArgAsScope(CI, 0);
diff --git a/llvm-spirv/lib/SPIRV/SPIRVToOCL.h b/llvm-spirv/lib/SPIRV/SPIRVToOCL.h
@@ -241,6 +241,9 @@ class SPIRVToOCLBase : public InstVisitor<SPIRVToOCLBase>,
   /// Transform relational builtin, e.g. __spirv_IsNan, to OpenCL builtin.
   void visitCallSPIRVRelational(CallInst *CI, Op OC);
 
+  /// Transform __spirv_ReadClockKHR to OpenCL builtin.
+  void visitCallSPIRVReadClockKHR(CallInst *CI);
+
   /// Conduct generic mutations for all atomic builtins
   virtual CallInst *mutateCommonAtomicArguments(CallInst *CI, Op OC) = 0;
 
diff --git a/llvm-spirv/test/extensions/KHR/SPV_KHR_shader_clock/shader_clock.cl b/llvm-spirv/test/extensions/KHR/SPV_KHR_shader_clock/shader_clock.cl
@@ -0,0 +1,50 @@
+// REQUIRES: spirv-dis
+// RUN: %clang_cc1 -triple spir-unknown-unknown -O1 -cl-std=CL2.0 -fdeclare-opencl-builtins -finclude-default-header -emit-llvm-bc %s -o %t.bc
+// RUN: llvm-spirv %t.bc --spirv-ext=+SPV_KHR_shader_clock -o %t.spv
+// RUN: spirv-dis %t.spv -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+// TODO: spirv-val %t.spv
+// RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+// RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+// RUN: llvm-spirv -r --spirv-target-env=SPV-IR %t.spv -o %t.rev.bc
+// RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-SPV-IR
+
+// CHECK-SPIRV: OpCapability ShaderClockKHR
+// CHECK-SPIRV: OpExtension "SPV_KHR_shader_clock"
+// CHECK-SPIRV-DAG: [[uint:%[a-z0-9_]+]] = OpTypeInt 32
+// CHECK-SPIRV-DAG: [[ulong:%[a-z0-9_]+]] = OpTypeInt 64
+// CHECK-SPIRV-DAG: [[v2uint:%[a-z0-9_]+]] = OpTypeVector [[uint]] 2
+// CHECK-SPIRV-DAG: [[uint_1:%[a-z0-9_]+]] = OpConstant [[uint]] 1
+// CHECK-SPIRV-DAG: [[uint_2:%[a-z0-9_]+]] = OpConstant [[uint]] 2
+// CHECK-SPIRV-DAG: [[uint_3:%[a-z0-9_]+]] = OpConstant [[uint]] 3
+// CHECK-SPIRV: OpReadClockKHR [[ulong]] [[uint_1]]
+// CHECK-SPIRV: OpReadClockKHR [[ulong]] [[uint_2]]
+// CHECK-SPIRV: OpReadClockKHR [[ulong]] [[uint_3]]
+// CHECK-SPIRV: OpReadClockKHR [[v2uint]] [[uint_1]]
+// CHECK-SPIRV: OpReadClockKHR [[v2uint]] [[uint_2]]
+// CHECK-SPIRV: OpReadClockKHR [[v2uint]] [[uint_3]]
+
+// CHECK-LLVM-LABEL: test_clocks
+// CHECK-LLVM: call spir_func i64 @_Z17clock_read_devicev()
+// CHECK-LLVM: call spir_func i64 @_Z21clock_read_work_groupv()
+// CHECK-LLVM: call spir_func i64 @_Z20clock_read_sub_groupv()
+// CHECK-LLVM: call spir_func <2 x i32> @_Z22clock_read_hilo_devicev()
+// CHECK-LLVM: call spir_func <2 x i32> @_Z26clock_read_hilo_work_groupv()
+// CHECK-LLVM: call spir_func <2 x i32> @_Z25clock_read_hilo_sub_groupv()
+
+// CHECK-SPV-IR-LABEL: test_clocks
+// CHECK-SPV-IR: call spir_func i64 @_Z27__spirv_ReadClockKHR_Rulongi(i32 1)
+// CHECK-SPV-IR: call spir_func i64 @_Z27__spirv_ReadClockKHR_Rulongi(i32 2)
+// CHECK-SPV-IR: call spir_func i64 @_Z27__spirv_ReadClockKHR_Rulongi(i32 3)
+// CHECK-SPV-IR: call spir_func <2 x i32> @_Z27__spirv_ReadClockKHR_Ruint2i(i32 1)
+// CHECK-SPV-IR: call spir_func <2 x i32> @_Z27__spirv_ReadClockKHR_Ruint2i(i32 2)
+// CHECK-SPV-IR: call spir_func <2 x i32> @_Z27__spirv_ReadClockKHR_Ruint2i(i32 3)
+
+kernel void test_clocks(global ulong *out64, global uint2 *outv2) {
+  out64[0] = clock_read_device();
+  out64[1] = clock_read_work_group();
+  out64[2] = clock_read_sub_group();
+
+  outv2[0] = clock_read_hilo_device();
+  outv2[1] = clock_read_hilo_work_group();
+  outv2[2] = clock_read_hilo_sub_group();
+}
