[Mlir-commits] [mlir] [mlir][SPIR-V] Convert math.cttz to spirv.GLFindILsb (PR #200455)
Arseniy Obolenskiy
llvmlistbot at llvm.org
Sat May 30 03:15:05 PDT 2026
https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/200455
>From b1ac94f067196289e925d75f654a04e22e2bdb58 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 29 May 2026 18:45:58 +0200
Subject: [PATCH 1/3] [mlir][SPIR-V] Convert math.cttz to GL FindILsb
---
.../Conversion/MathToSPIRV/MathToSPIRV.cpp | 42 ++++++++++++++++++-
.../MathToSPIRV/math-to-gl-spirv.mlir | 22 ++++++++++
2 files changed, 63 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index cf5ea9716b9f7..6eeb8f842a87c 100644
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -242,6 +242,45 @@ struct CountLeadingZerosPattern final
}
};
+/// Converts math.cttz to GL FindILsb. GL FindILsb returns -1 for a zero
+/// input while math.cttz must return the bitwidth, so the zero case is
+/// patched up with a select.
+struct CountTrailingZerosPattern final
+ : public OpConversionPattern<math::CountTrailingZerosOp> {
+ using Base::Base;
+
+ LogicalResult
+ matchAndRewrite(math::CountTrailingZerosOp countOp, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (LogicalResult res = checkSourceOpTypes(rewriter, countOp); failed(res))
+ return res;
+
+ Type type = getTypeConverter()->convertType(countOp.getType());
+ if (!type)
+ return failure();
+
+ unsigned bitwidth = 0;
+ if (isa<IntegerType>(type))
+ bitwidth = type.getIntOrFloatBitWidth();
+ if (auto vectorType = dyn_cast<VectorType>(type))
+ bitwidth = vectorType.getElementTypeBitWidth();
+ if (bitwidth != 32)
+ return failure();
+
+ Location loc = countOp.getLoc();
+ Value input = adaptor.getOperand();
+ Value val0 = getScalarOrVectorI32Constant(type, 0, rewriter, loc);
+ Value valBitwidth =
+ getScalarOrVectorI32Constant(type, bitwidth, rewriter, loc);
+
+ Value lsb = spirv::GLFindILsbOp::create(rewriter, loc, input);
+ Value isZero = spirv::IEqualOp::create(rewriter, loc, input, val0);
+ rewriter.replaceOpWithNewOp<spirv::SelectOp>(countOp, isZero, valBitwidth,
+ lsb);
+ return success();
+ }
+};
+
/// Converts math.expm1 to SPIR-V ops.
///
/// SPIR-V does not have a direct operations for exp(x)-1. Explicitly lower to
@@ -530,7 +569,8 @@ void populateMathToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
// GLSL patterns
patterns.add<
- CountLeadingZerosPattern, Log1pOpPattern<spirv::GLLogOp>, Log10OpPattern,
+ CountLeadingZerosPattern, CountTrailingZerosPattern,
+ Log1pOpPattern<spirv::GLLogOp>, Log10OpPattern,
ExpM1OpPattern<spirv::GLExpOp>, PowFOpPattern, RoundOpPattern,
CheckedElementwiseOpPattern<math::AbsFOp, spirv::GLFAbsOp>,
CheckedElementwiseOpPattern<math::AbsIOp, spirv::GLSAbsOp>,
diff --git a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
index 08d7822d04cc1..94c285ddb0ccc 100644
--- a/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
+++ b/mlir/test/Conversion/MathToSPIRV/math-to-gl-spirv.mlir
@@ -183,6 +183,28 @@ func.func @ctlz_vector2(%val: vector<2xi32>) -> vector<2xi32> {
return %0 : vector<2xi32>
}
+// CHECK-LABEL: @cttz_scalar
+// CHECK-SAME: (%[[VAL:.+]]: i32)
+func.func @cttz_scalar(%val: i32) -> i32 {
+ // CHECK-DAG: %[[V0:.+]] = spirv.Constant 0 : i32
+ // CHECK-DAG: %[[V32:.+]] = spirv.Constant 32 : i32
+ // CHECK: %[[LSB:.+]] = spirv.GL.FindILsb %[[VAL]] : i32
+ // CHECK: %[[CMP:.+]] = spirv.IEqual %[[VAL]], %[[V0]] : i32
+ // CHECK: %[[R:.+]] = spirv.Select %[[CMP]], %[[V32]], %[[LSB]] : i1, i32
+ // CHECK: return %[[R]]
+ %0 = math.cttz %val : i32
+ return %0 : i32
+}
+
+// CHECK-LABEL: @cttz_vector
+func.func @cttz_vector(%val: vector<2xi32>) -> vector<2xi32> {
+ // CHECK: spirv.GL.FindILsb
+ // CHECK: spirv.IEqual
+ // CHECK: spirv.Select
+ %0 = math.cttz %val : vector<2xi32>
+ return %0 : vector<2xi32>
+}
+
// Dynamic exponent: exp(y * log(x)); yields NaN for x<0.
// CHECK-LABEL: @powf_scalar
// CHECK-SAME: (%[[LHS:.+]]: f32, %[[RHS:.+]]: f32)
>From 4df61e2e79740d848677c0d05377a323730eab10 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Fri, 29 May 2026 19:10:24 +0200
Subject: [PATCH 2/3] fmt
---
.../Conversion/MathToSPIRV/MathToSPIRV.cpp | 58 +++++++++----------
1 file changed, 29 insertions(+), 29 deletions(-)
diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index 6eeb8f842a87c..8d40671fab0f1 100644
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -568,35 +568,35 @@ void populateMathToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
typeConverter, patterns.getContext());
// GLSL patterns
- patterns.add<
- CountLeadingZerosPattern, CountTrailingZerosPattern,
- Log1pOpPattern<spirv::GLLogOp>, Log10OpPattern,
- ExpM1OpPattern<spirv::GLExpOp>, PowFOpPattern, RoundOpPattern,
- CheckedElementwiseOpPattern<math::AbsFOp, spirv::GLFAbsOp>,
- CheckedElementwiseOpPattern<math::AbsIOp, spirv::GLSAbsOp>,
- CheckedElementwiseOpPattern<math::AtanOp, spirv::GLAtanOp>,
- CheckedElementwiseOpPattern<math::CeilOp, spirv::GLCeilOp>,
- CheckedElementwiseOpPattern<math::CosOp, spirv::GLCosOp>,
- CheckedElementwiseOpPattern<math::ExpOp, spirv::GLExpOp>,
- CheckedElementwiseOpPattern<math::Exp2Op, spirv::GLExp2Op>,
- CheckedElementwiseOpPattern<math::FloorOp, spirv::GLFloorOp>,
- CheckedElementwiseOpPattern<math::FmaOp, spirv::GLFmaOp>,
- CheckedElementwiseOpPattern<math::LogOp, spirv::GLLogOp>,
- CheckedElementwiseOpPattern<math::Log2Op, spirv::GLLog2Op>,
- CheckedElementwiseOpPattern<math::RoundEvenOp, spirv::GLRoundEvenOp>,
- CheckedElementwiseOpPattern<math::RsqrtOp, spirv::GLInverseSqrtOp>,
- CheckedElementwiseOpPattern<math::SinOp, spirv::GLSinOp>,
- CheckedElementwiseOpPattern<math::SqrtOp, spirv::GLSqrtOp>,
- CheckedElementwiseOpPattern<math::TanhOp, spirv::GLTanhOp>,
- CheckedElementwiseOpPattern<math::TanOp, spirv::GLTanOp>,
- CheckedElementwiseOpPattern<math::AsinOp, spirv::GLAsinOp>,
- CheckedElementwiseOpPattern<math::AcosOp, spirv::GLAcosOp>,
- CheckedElementwiseOpPattern<math::SinhOp, spirv::GLSinhOp>,
- CheckedElementwiseOpPattern<math::CoshOp, spirv::GLCoshOp>,
- CheckedElementwiseOpPattern<math::AsinhOp, spirv::GLAsinhOp>,
- CheckedElementwiseOpPattern<math::AcoshOp, spirv::GLAcoshOp>,
- CheckedElementwiseOpPattern<math::AtanhOp, spirv::GLAtanhOp>>(
- typeConverter, patterns.getContext());
+ patterns
+ .add<CountLeadingZerosPattern, CountTrailingZerosPattern,
+ Log1pOpPattern<spirv::GLLogOp>, Log10OpPattern,
+ ExpM1OpPattern<spirv::GLExpOp>, PowFOpPattern, RoundOpPattern,
+ CheckedElementwiseOpPattern<math::AbsFOp, spirv::GLFAbsOp>,
+ CheckedElementwiseOpPattern<math::AbsIOp, spirv::GLSAbsOp>,
+ CheckedElementwiseOpPattern<math::AtanOp, spirv::GLAtanOp>,
+ CheckedElementwiseOpPattern<math::CeilOp, spirv::GLCeilOp>,
+ CheckedElementwiseOpPattern<math::CosOp, spirv::GLCosOp>,
+ CheckedElementwiseOpPattern<math::ExpOp, spirv::GLExpOp>,
+ CheckedElementwiseOpPattern<math::Exp2Op, spirv::GLExp2Op>,
+ CheckedElementwiseOpPattern<math::FloorOp, spirv::GLFloorOp>,
+ CheckedElementwiseOpPattern<math::FmaOp, spirv::GLFmaOp>,
+ CheckedElementwiseOpPattern<math::LogOp, spirv::GLLogOp>,
+ CheckedElementwiseOpPattern<math::Log2Op, spirv::GLLog2Op>,
+ CheckedElementwiseOpPattern<math::RoundEvenOp, spirv::GLRoundEvenOp>,
+ CheckedElementwiseOpPattern<math::RsqrtOp, spirv::GLInverseSqrtOp>,
+ CheckedElementwiseOpPattern<math::SinOp, spirv::GLSinOp>,
+ CheckedElementwiseOpPattern<math::SqrtOp, spirv::GLSqrtOp>,
+ CheckedElementwiseOpPattern<math::TanhOp, spirv::GLTanhOp>,
+ CheckedElementwiseOpPattern<math::TanOp, spirv::GLTanOp>,
+ CheckedElementwiseOpPattern<math::AsinOp, spirv::GLAsinOp>,
+ CheckedElementwiseOpPattern<math::AcosOp, spirv::GLAcosOp>,
+ CheckedElementwiseOpPattern<math::SinhOp, spirv::GLSinhOp>,
+ CheckedElementwiseOpPattern<math::CoshOp, spirv::GLCoshOp>,
+ CheckedElementwiseOpPattern<math::AsinhOp, spirv::GLAsinhOp>,
+ CheckedElementwiseOpPattern<math::AcoshOp, spirv::GLAcoshOp>,
+ CheckedElementwiseOpPattern<math::AtanhOp, spirv::GLAtanhOp>>(
+ typeConverter, patterns.getContext());
// OpenCL patterns
patterns.add<
>From 808840dd2de8b565f28870f51eb3384002c38c6a Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Sat, 30 May 2026 12:01:12 +0200
Subject: [PATCH 3/3] Address comment
---
mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
index 8d40671fab0f1..3d7f76b90e7c5 100644
--- a/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
+++ b/mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp
@@ -262,7 +262,7 @@ struct CountTrailingZerosPattern final
unsigned bitwidth = 0;
if (isa<IntegerType>(type))
bitwidth = type.getIntOrFloatBitWidth();
- if (auto vectorType = dyn_cast<VectorType>(type))
+ else if (auto vectorType = dyn_cast<VectorType>(type))
bitwidth = vectorType.getElementTypeBitWidth();
if (bitwidth != 32)
return failure();
More information about the Mlir-commits
mailing list