[Mlir-commits] [mlir] e6c20ad - [mlir][gpu][math] Fix assertion in OpToFuncCallLowering on vector results (#215317)

[llvmlistbot at llvm.org]

Author: SeongJaePark
Date: 2026-08-31T11:08:22+02:00
New Revision: e6c20ada052d34166da21bdb40f79a4e177bc115

URL: https://github.com/llvm/llvm-project/commit/e6c20ada052d34166da21bdb40f79a4e177bc115
DIFF: https://github.com/llvm/llvm-project/commit/e6c20ada052d34166da21bdb40f79a4e177bc115.diff

LOG: [mlir][gpu][math] Fix assertion in OpToFuncCallLowering on vector results (#215317)

`OpToFuncCallLowering` asserts that operand and result types match, with
an escape hatch for scalar `i1` results. Math vector FP classification ops
return `vector<...xi1>`, which the escape hatch does not cover, so they trigger
the assertion.

Reject non-scalar results before the assertion, allowing
`ScalarizeVectorOpLowering` to lower them element-wise. This preserves
the existing behavior in non-assert builds, where `OpToFuncCallLowering`
already failed to match vector types. Add regression coverage for vector and
scalar FP classification ops.

Fixes #210855

Assisted-by: Claude (Anthropic)

Added: 
    mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir

Modified: 
    mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
index cb9b6da071839..ea6d7c9f2da1d 100644
--- a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
+++ b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
@@ -72,7 +72,15 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
         std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,
         "expected single result op");
 
-    bool isResultBool = op->getResultTypes().front().isInteger(1);
+    // This pattern only handles scalar ops. Ops with shaped (e.g. vector)
+    // result types, such as `math.isinf` on `vector<Nxf32>`, are expected to be
+    // scalarized first by `ScalarizeVectorOpLowering`, which is co-registered
+    // for these ops; bail out so that pattern can take over.
+    Type opResultType = op->getResultTypes().front();
+    if (!opResultType.isIntOrIndexOrFloat())
+      return rewriter.notifyMatchFailure(op, "expected scalar result type");
+
+    bool isResultBool = opResultType.isInteger(1);
     if constexpr (!std::is_base_of<OpTrait::SameOperandsAndResultType<SourceOp>,
                                    SourceOp>::value) {
       assert(op->getNumOperands() > 0 &&

diff  --git a/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir b/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir
new file mode 100644
index 0000000000000..ac14ab15fd80c
--- /dev/null
+++ b/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -convert-math-to-nvvm | FileCheck %s
+
+// Classification ops return a bool, so their operand and result types 
diff er.
+// On shaped operands `OpToFuncCallLowering` bails out and `ScalarizeVectorOpLowering`
+// unrolls them element-wise, lowering each element to a libdevice call.
+
+// CHECK-LABEL:   func.func @fpclass_vector(
+// CHECK-SAME:                              %[[ARG:.*]]: vector<2xf32>)
+func.func @fpclass_vector(%arg: vector<2xf32>) -> (vector<2xi1>, vector<2xi1>, vector<2xi1>) {
+  // CHECK-COUNT-2: llvm.call @__nv_isinff({{.*}}) : (f32) -> i32
+  %inf = math.isinf %arg : vector<2xf32>
+  // CHECK-COUNT-2: llvm.call @__nv_finitef({{.*}}) : (f32) -> i32
+  %finite = math.isfinite %arg : vector<2xf32>
+  // CHECK-COUNT-2: llvm.call @__nv_isnanf({{.*}}) : (f32) -> i32
+  %nan = math.isnan %arg : vector<2xf32>
+  return %inf, %finite, %nan : vector<2xi1>, vector<2xi1>, vector<2xi1>
+}
+
+// CHECK-LABEL:   func.func @fpclass_scalar(
+func.func @fpclass_scalar(%arg: f32) -> (i1, i1, i1) {
+  // CHECK: %[[INF:.*]] = llvm.call @__nv_isinff({{.*}}) : (f32) -> i32
+  // CHECK: llvm.icmp "ne" %[[INF]], {{.*}} : i32
+  %inf = math.isinf %arg : f32
+  // CHECK: llvm.call @__nv_finitef({{.*}}) : (f32) -> i32
+  %finite = math.isfinite %arg : f32
+  // CHECK: llvm.call @__nv_isnanf({{.*}}) : (f32) -> i32
+  %nan = math.isnan %arg : f32
+  return %inf, %finite, %nan : i1, i1, i1
+}