[Mlir-commits] [mlir] [mlir][spirv] Support OpenCL.std clz in the SPIR-V dialect (PR #195317)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri May 1 11:31:25 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Levin Dabhi (levindabhi)
<details>
<summary>Changes</summary>
Add support for the OpenCL.std `clz` extended instruction in the MLIR SPIR-V dialect.
Add `spirv.CL.clz` to the SPIR-V OpenCL extended instruction ops and lower `math.ctlz` to it for targets with Kernel capability. Shader targets keep the existing GLSL-based fallback implemented via `spirv.GL.FindUMsb`.
Previously, `math.ctlz` was lowered through the GLSL path using `spirv.GL.FindUMsb` plus additional SPIR-V ops. That worked for Shader targets, but failed legalization for OpenCL/Kernel targets where Shader capability is not supported.
---
Full diff: https://github.com/llvm/llvm-project/pull/195317.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td (+23)
- (modified) mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp (+42-33)
- (modified) mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir (+9)
- (modified) mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir (+50)
- (modified) mlir/test/Target/SPIRV/ocl-ops.mlir (+2)
``````````diff
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td
index d36245d5ad6b7..71f9c9579db81 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCLOps.td
@@ -783,6 +783,29 @@ def SPIRV_CLSAbsOp : SPIRV_CLUnaryArithmeticOp<"s_abs", 141, SPIRV_Integer> {
// -----
+def SPIRV_CLClzOp : SPIRV_CLUnaryArithmeticOp<"clz", 151, SPIRV_Integer> {
+ let summary = "Count leading zeros in operand";
+
+ let description = [{
+ Returns the number of most-significant 0-bits in x.
+
+ Result Type and x must be integer or vector(2,3,4,8,16) of
+ integer values.
+
+ All of the operands, including the Result Type operand, must be of the
+ same type.
+
+ #### Example:
+
+ ```mlir
+ %2 = spirv.CL.clz %0 : i32
+ %3 = spirv.CL.clz %1 : vector<3xi16>
+ ```
+ }];
+}
+
+// -----
+
def SPIRV_CLSMaxOp : SPIRV_CLBinaryArithmeticOp<"s_max", 156, SPIRV_Integer> {
let summary = "Return maximum of two signed integer operands";
diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index e4b5da7a5ea92..ea6be76373573 100644
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -183,9 +183,12 @@ struct CopySignPattern final : public OpConversionPattern<math::CopySignOp> {
/// Converts math.ctlz to SPIR-V ops.
///
-/// SPIR-V does not have a direct operations for counting leading zeros. If
-/// Shader capability is supported, we can leverage GL FindUMsb to calculate
-/// it.
+/// OpenCL targets lower math.ctlz directly to OpenCL.std clz via the generic
+/// elementwise pattern. This pattern handles the shader fallback.
+///
+/// SPIR-V does not have a direct operations for counting leading zeros for
+/// glsl. If Shader capability is supported, we can leverage GL FindUMsb to
+/// calculate it.
struct CountLeadingZerosPattern final
: public OpConversionPattern<math::CountLeadingZerosOp> {
using Base::Base;
@@ -200,7 +203,11 @@ struct CountLeadingZerosPattern final
if (!type)
return failure();
- // We can only support 32-bit integer types for now.
+ auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
+ if (!typeConverter.getTargetEnv().allows(spirv::Capability::Shader))
+ return rewriter.notifyMatchFailure(countOp, "requires Shader capability");
+
+ // The GL FindUMsb fallback only supports 32-bit integer types for now.
unsigned bitwidth = 0;
if (isa<IntegerType>(type))
bitwidth = type.getIntOrFloatBitWidth();
@@ -533,35 +540,37 @@ void populateMathToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
typeConverter, patterns.getContext());
// OpenCL patterns
- patterns.add<Log1pOpPattern<spirv::CLLogOp>, ExpM1OpPattern<spirv::CLExpOp>,
- Log2Log10OpPattern<math::Log2Op, spirv::CLLogOp>,
- Log2Log10OpPattern<math::Log10Op, spirv::CLLogOp>,
- CheckedElementwiseOpPattern<math::AbsFOp, spirv::CLFAbsOp>,
- CheckedElementwiseOpPattern<math::AbsIOp, spirv::CLSAbsOp>,
- CheckedElementwiseOpPattern<math::AtanOp, spirv::CLAtanOp>,
- CheckedElementwiseOpPattern<math::Atan2Op, spirv::CLAtan2Op>,
- CheckedElementwiseOpPattern<math::CeilOp, spirv::CLCeilOp>,
- CheckedElementwiseOpPattern<math::CosOp, spirv::CLCosOp>,
- CheckedElementwiseOpPattern<math::ErfOp, spirv::CLErfOp>,
- CheckedElementwiseOpPattern<math::ExpOp, spirv::CLExpOp>,
- CheckedElementwiseOpPattern<math::FloorOp, spirv::CLFloorOp>,
- CheckedElementwiseOpPattern<math::FmaOp, spirv::CLFmaOp>,
- CheckedElementwiseOpPattern<math::LogOp, spirv::CLLogOp>,
- CheckedElementwiseOpPattern<math::PowFOp, spirv::CLPowOp>,
- CheckedElementwiseOpPattern<math::RoundEvenOp, spirv::CLRintOp>,
- CheckedElementwiseOpPattern<math::RoundOp, spirv::CLRoundOp>,
- CheckedElementwiseOpPattern<math::RsqrtOp, spirv::CLRsqrtOp>,
- CheckedElementwiseOpPattern<math::SinOp, spirv::CLSinOp>,
- CheckedElementwiseOpPattern<math::SqrtOp, spirv::CLSqrtOp>,
- CheckedElementwiseOpPattern<math::TanhOp, spirv::CLTanhOp>,
- CheckedElementwiseOpPattern<math::TanOp, spirv::CLTanOp>,
- CheckedElementwiseOpPattern<math::AsinOp, spirv::CLAsinOp>,
- CheckedElementwiseOpPattern<math::AcosOp, spirv::CLAcosOp>,
- CheckedElementwiseOpPattern<math::SinhOp, spirv::CLSinhOp>,
- CheckedElementwiseOpPattern<math::CoshOp, spirv::CLCoshOp>,
- CheckedElementwiseOpPattern<math::AsinhOp, spirv::CLAsinhOp>,
- CheckedElementwiseOpPattern<math::AcoshOp, spirv::CLAcoshOp>,
- CheckedElementwiseOpPattern<math::AtanhOp, spirv::CLAtanhOp>>(
+ patterns.add<
+ Log1pOpPattern<spirv::CLLogOp>, ExpM1OpPattern<spirv::CLExpOp>,
+ Log2Log10OpPattern<math::Log2Op, spirv::CLLogOp>,
+ Log2Log10OpPattern<math::Log10Op, spirv::CLLogOp>,
+ CheckedElementwiseOpPattern<math::AbsFOp, spirv::CLFAbsOp>,
+ CheckedElementwiseOpPattern<math::AbsIOp, spirv::CLSAbsOp>,
+ CheckedElementwiseOpPattern<math::CountLeadingZerosOp, spirv::CLClzOp>,
+ CheckedElementwiseOpPattern<math::AtanOp, spirv::CLAtanOp>,
+ CheckedElementwiseOpPattern<math::Atan2Op, spirv::CLAtan2Op>,
+ CheckedElementwiseOpPattern<math::CeilOp, spirv::CLCeilOp>,
+ CheckedElementwiseOpPattern<math::CosOp, spirv::CLCosOp>,
+ CheckedElementwiseOpPattern<math::ErfOp, spirv::CLErfOp>,
+ CheckedElementwiseOpPattern<math::ExpOp, spirv::CLExpOp>,
+ CheckedElementwiseOpPattern<math::FloorOp, spirv::CLFloorOp>,
+ CheckedElementwiseOpPattern<math::FmaOp, spirv::CLFmaOp>,
+ CheckedElementwiseOpPattern<math::LogOp, spirv::CLLogOp>,
+ CheckedElementwiseOpPattern<math::PowFOp, spirv::CLPowOp>,
+ CheckedElementwiseOpPattern<math::RoundEvenOp, spirv::CLRintOp>,
+ CheckedElementwiseOpPattern<math::RoundOp, spirv::CLRoundOp>,
+ CheckedElementwiseOpPattern<math::RsqrtOp, spirv::CLRsqrtOp>,
+ CheckedElementwiseOpPattern<math::SinOp, spirv::CLSinOp>,
+ CheckedElementwiseOpPattern<math::SqrtOp, spirv::CLSqrtOp>,
+ CheckedElementwiseOpPattern<math::TanhOp, spirv::CLTanhOp>,
+ CheckedElementwiseOpPattern<math::TanOp, spirv::CLTanOp>,
+ CheckedElementwiseOpPattern<math::AsinOp, spirv::CLAsinOp>,
+ CheckedElementwiseOpPattern<math::AcosOp, spirv::CLAcosOp>,
+ CheckedElementwiseOpPattern<math::SinhOp, spirv::CLSinhOp>,
+ CheckedElementwiseOpPattern<math::CoshOp, spirv::CLCoshOp>,
+ CheckedElementwiseOpPattern<math::AsinhOp, spirv::CLAsinhOp>,
+ CheckedElementwiseOpPattern<math::AcoshOp, spirv::CLAcoshOp>,
+ CheckedElementwiseOpPattern<math::AtanhOp, spirv::CLAtanhOp>>(
typeConverter, patterns.getContext());
}
diff --git a/mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir b/mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir
index 56a0d4dafec8c..dae1b43402718 100644
--- a/mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir
+++ b/mlir/test/Conversion/MathToSPIRV/math-to-opencl-spirv.mlir
@@ -159,6 +159,15 @@ func.func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
func.func @int_unary(%arg0: i32) {
// CHECK: spirv.CL.s_abs %{{.*}}
%0 = math.absi %arg0 : i32
+ // CHECK: spirv.CL.clz %{{.*}} : i32
+ %1 = math.ctlz %arg0 : i32
+ return
+}
+
+// CHECK-LABEL: @int_unary_vector
+func.func @int_unary_vector(%arg0: vector<2xi32>) {
+ // CHECK: spirv.CL.clz %{{.*}} : vector<2xi32>
+ %0 = math.ctlz %arg0 : vector<2xi32>
return
}
diff --git a/mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir b/mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir
index 8d81cb42030a5..d6d3c53f23356 100644
--- a/mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/ocl-ops.mlir
@@ -168,6 +168,56 @@ func.func @sabs(%arg0 : i32) -> () {
// -----
+//===----------------------------------------------------------------------===//
+// spirv.CL.clz
+//===----------------------------------------------------------------------===//
+
+func.func @clz(%arg0 : i32) -> () {
+ // CHECK: spirv.CL.clz {{%.*}} : i32
+ %2 = spirv.CL.clz %arg0 : i32
+ return
+}
+
+func.func @clzvec(%arg0 : vector<3xi16>) -> () {
+ // CHECK: spirv.CL.clz {{%.*}} : vector<3xi16>
+ %2 = spirv.CL.clz %arg0 : vector<3xi16>
+ return
+}
+
+// -----
+
+func.func @clz(%arg0 : f32) -> () {
+ // expected-error @+1 {{op operand #0 must be 8/16/32/64-bit integer or fixed-length vector of 8/16/32/64-bit integer values}}
+ %2 = spirv.CL.clz %arg0 : f32
+ return
+}
+
+// -----
+
+func.func @clz(%arg0 : vector<5xi32>) -> () {
+ // expected-error @+1 {{op operand #0 must be 8/16/32/64-bit integer or fixed-length vector of 8/16/32/64-bit integer values of length 2/3/4}}
+ %2 = spirv.CL.clz %arg0 : vector<5xi32>
+ return
+}
+
+// -----
+
+func.func @clz(%arg0 : i32, %arg1 : i32) -> () {
+ // expected-error @+1 {{expected ':'}}
+ %2 = spirv.CL.clz %arg0, %arg1 : i32
+ return
+}
+
+// -----
+
+func.func @clz(%arg0 : i32) -> () {
+ // expected-error @+1 {{expected non-function type}}
+ %2 = spirv.CL.clz %arg0 :
+ return
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.CL.fma
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/SPIRV/ocl-ops.mlir b/mlir/test/Target/SPIRV/ocl-ops.mlir
index 14e45518502de..7a4abbd9dd344 100644
--- a/mlir/test/Target/SPIRV/ocl-ops.mlir
+++ b/mlir/test/Target/SPIRV/ocl-ops.mlir
@@ -35,6 +35,8 @@ spirv.module Physical64 OpenCL requires #spirv.vce<v1.0, [Kernel, Addresses, Vec
spirv.func @integer_insts(%arg0 : i32) "None" {
// CHECK: {{%.*}} = spirv.CL.s_abs {{%.*}} : i32
%0 = spirv.CL.s_abs %arg0 : i32
+ // CHECK: {{%.*}} = spirv.CL.clz {{%.*}} : i32
+ %1 = spirv.CL.clz %arg0 : i32
spirv.Return
}
``````````
</details>
https://github.com/llvm/llvm-project/pull/195317
More information about the Mlir-commits
mailing list