[mlir][gpu][math] Fix assertion in OpToFuncCallLowering on vector results

Project: llvm/llvm-project | Language: C++ | Concepts: error-handling, refactoring, testing-strategy | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/215317

The problem

The OpToFuncCallLowering pattern in MLIR's GPU math conversion pipeline asserted that operand and result types must match, with a special case for scalar boolean (i1) results. However, some math ops (like floating point classification) return vector<...xi1> when applied to vectors, which was not handled by the escape hatch. This caused assertion failures when lowering vector math classification ops (e.g., math.isinf on vector<2xf32>).

Why it's non-trivial

A naive fix might try to relax the assertion to accept vector boolean types, but this would allow the lowering path to proceed with mismatched types, producing incorrect IR or hiding real mismatches. The right fix must ensure only scalar ops are handled here and defer vector cases to the appropriate scalarization pattern, preserving correct lowering and error reporting.

The approach

The patch changes the matching logic in OpToFuncCallLowering to explicitly reject non-scalar result types before reaching the assertion, allowing the scalarization pattern to handle vectors. It also adds targeted regression tests.

Key code change:

-    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);

• What it does:
  • Checks if the result type is a scalar (Int, Index, or Float).
  • If not (i.e., it's a vector or other shaped type), it returns a match failure with a clear reason.
  • This prevents the assertion from firing on vector types and lets the scalarization lowering pattern handle them.

Test coverage added:

func.func @fpclass_vector(%arg: vector<2xf32>) -> (vector<2xi1>, vector<2xi1>, vector<2xi1>) {
  %inf = math.isinf %arg : vector<2xf32>
  %finite = math.isfinite %arg : vector<2xf32>
  %nan = math.isnan %arg : vector<2xf32>
  return %inf, %finite, %nan : vector<2xi1>, vector<2xi1>, vector<2xi1>
}

func.func @fpclass_scalar(%arg: f32) -> (i1, i1, i1) {
  %inf = math.isinf %arg : f32
  %finite = math.isfinite %arg : f32
  %nan = math.isnan %arg : f32
  return %inf, %finite, %nan : i1, i1, i1
}

• What it does:
  • Verifies that scalar ops are lowered directly, while vector ops are handled element-wise by the scalarization pattern, calling the appropriate libdevice functions per element.

From the review thread

Also please trim the commit description, it should be straight to the point, as it's a straightforward fix.

Done — I've shortend the description. Thanks!

Takeaways

• Always explicitly check and reject unsupported types early in a lowering or transformation pattern to avoid hidden assertion failures and improve error reporting.
• Defer handling of shaped (e.g., vector) types to specialized scalarization or unrolling patterns, rather than forcing all logic into a scalar-oriented lowering.
• Add regression tests covering both the problematic and the correct cases to ensure the fix is robust and future changes do not reintroduce the bug.