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

[llvmlistbot at llvm.org]

https://github.com/SeongjaeP created https://github.com/llvm/llvm-project/pull/215317

`OpToFuncCallLowering` lowers scalar math ops to target runtime function calls
such as libdevice, OCML, and OCL functions. For ops without the
`SameOperandsAndResultType` trait it asserts that the operand and result types
match, with an escape hatch for bool results:

```c++
bool isResultBool = op->getResultTypes().front().isInteger(1);
```

That escape hatch only fires for scalars: `Type::isInteger(1)` is false for
`vector<2xi1>`, which is a `VectorType`, not an `IntegerType`. So the FP
classification ops (`math.isinf`, `math.isfinite`, `math.isnan`), whose operand
and result types genuinely differ, hit the assertion as soon as they are applied
to a vector:

```
$ mlir-opt repro.mlir -convert-math-to-nvvm
mlir-opt: OpToFuncCallLowering.h:80: ... Assertion
  `(op->getResultTypes().front() == op->getOperand(0).getType() || isResultBool)
   && "expected op with same operand and result types"' failed.
```

## Fix

Bail out on non-scalar results before the assertion and let the co-registered
`ScalarizeVectorOpLowering` handle them:

```c++
Type opResultType = op->getResultTypes().front();
if (!opResultType.isIntOrIndexOrFloat())
  return rewriter.notifyMatchFailure(op, "expected scalar result type");

bool isResultBool = opResultType.isInteger(1);
```

For builds without assertions, this does not change the generated IR. The
pattern already returned `failure()` for vector types when `getFunctionName()`
returned an empty name, allowing the co-registered `ScalarizeVectorOpLowering`
to handle the operation. This check moves that decision before the assertion,
so assertion-enabled and release builds follow the same lowering path, and the
reported match failure states the actual reason. With
`-debug-only=dialect-conversion`:

```
Legalizing operation : 'math.isinf' {
  * Pattern : 'math.isinf -> ()' {
    ** Failure : expected scalar result type
  } -> FAILURE : pattern failed to match
  * Pattern : 'math.isinf -> ()' {
    ** Insert  : 'llvm.extractelement' ... 'math.isinf' ... 'llvm.insertelement'
    ** Replace : 'math.isinf'
```

### Why not relax `isResultBool` instead

The issue suggested `getElementTypeOrSelf(...).isInteger(1)`. I built that
variant and confirmed it produces IR identical to this patch today — but only
because the pattern still exits via the `funcName.empty()` path further down,
without reporting a match failure reason. By then `isResultBool` has been used
to decide `resultType = i32`, which is not meaningful for a `vector<2xi1>`
result; the scalar-only tail of the pattern (returning the result as `i32` per
the ABI and comparing it with `icmp ne %r, 0 : i32`) would emit wrong IR if it
were ever reached, e.g. if `getFunctionName()` later learned about vector types.
It would also weaken the assertion, which would then pass for operand/result
types that genuinely differ.

### Other consumers of `OpToFuncCallLowering`

The issue asked whether the other consumers need auditing.
`OpToFuncCallLowering` is instantiated in three places, all of which
co-register `ScalarizeVectorOpLowering` for the same operations:

- `MathToNVVM.cpp`
- `MathToROCDL.cpp`
- `MathToXeVM.cpp`

A `vectorizable`-trait check does not seem necessary here: this pattern only
needs to reject non-scalar results before entering its scalar-call lowering.
Whether an operation may be unrolled element-wise is determined when the caller
registers `ScalarizeVectorOpLowering`. `impl::scalarizeVectorOp` handles both
`VectorType` and `LLVM::LLVMArrayType` operands.

The new test covers both vector classification — lowered to per-element
`__nv_isinff`, `__nv_finitef`, and `__nv_isnanf` calls — and unchanged scalar
lowering.

Fixes #210855


>From 46cbbf6935298c263627a7ae29465b13972087b8 Mon Sep 17 00:00:00 2001
From: SeongjaeP <psjj960507 at gmail.com>
Date: Mon, 10 Aug 2026 13:59:31 +0000
Subject: [PATCH] [mlir][gpu][math] Fix assertion in OpToFuncCallLowering on
 vector results

---
 .../GPUCommon/OpToFuncCallLowering.h          | 10 ++++++-
 .../Conversion/MathToNVVM/math-to-nvvm.mlir   | 29 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir

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 differ.
+// 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
+}