[mlir][math] Fix intrinsic conversions to LLVM for 0D-vector types

[AGindinson]

vector<t> types are not compatible with the LLVM type system – with the current approach employed within LLVMTypeConverter, they must be explicitly converted into vector<1xt> when lowering. Employ this rule within the conversion patterns for intrinsics that are handled directly within MathToLLVM: math.ctlz .cttz, .absi, .expm1, .log1p, .rsqrt, .isnan, .isfinite.

This change does not cover/test patterns that are based off VectorConvertToLLVMPattern template from LLVMCommon/VectorPattern.h.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Artem Gindinson (AGindinson)

Changes

vector<t> types are not compatible with the LLVM type system, and must be explicitly converted into vector<1xt> when lowering. Employ this rule within the conversion pattern for math.ctlz, .cttz and .absi intrinsics.

---

Full diff: https://github.com/llvm/llvm-project/pull/141020.diff

2 Files Affected:

• (modified) mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp (+10-1)
• (modified) mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir (+33)

diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 97da96afac4cd..19cd960b15294 100644
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -84,6 +84,15 @@ struct IntOpWithFlagLowering : public ConvertOpToLLVMPattern<MathOp> {
 
     auto loc = op.getLoc();
     auto resultType = op.getResult().getType();
+    const auto &typeConverter = *this->getTypeConverter();
+    if (!LLVM::isCompatibleType(resultType)) {
+      resultType = typeConverter.convertType(resultType);
+      if (!resultType)
+        return failure();
+    }
+    if (operandType != resultType)
+      return rewriter.notifyMatchFailure(
+          op, "compatible result type doesn't match operand type");
 
     if (!isa<LLVM::LLVMArrayType>(operandType)) {
       rewriter.replaceOpWithNewOp<LLVMOp>(op, resultType, adaptor.getOperand(),
@@ -96,7 +105,7 @@ struct IntOpWithFlagLowering : public ConvertOpToLLVMPattern<MathOp> {
       return failure();
 
     return LLVM::detail::handleMultidimensionalVectors(
-        op.getOperation(), adaptor.getOperands(), *this->getTypeConverter(),
+        op.getOperation(), adaptor.getOperands(), typeConverter,
         [&](Type llvm1DVectorTy, ValueRange operands) {
           return rewriter.create<LLVMOp>(loc, llvm1DVectorTy, operands[0],
                                          false);
diff --git a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
index 974743a55932b..73325a3fd913e 100644
--- a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
+++ b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
@@ -19,6 +19,8 @@ func.func @ops(%arg0: f32, %arg1: f32, %arg2: i32, %arg3: i32, %arg4: f64) {
 
 // -----
 
+// CHECK-LABEL: func @absi(
+// CHECK-SAME: i32
 func.func @absi(%arg0: i32) -> i32 {
   // CHECK: = "llvm.intr.abs"(%{{.*}}) <{is_int_min_poison = false}> : (i32) -> i32
   %0 = math.absi %arg0 : i32
@@ -27,6 +29,17 @@ func.func @absi(%arg0: i32) -> i32 {
 
 // -----
 
+// CHECK-LABEL: func @absi_0d_vec(
+// CHECK-SAME: i32
+func.func @absi_0d_vec(%arg0 : vector<i32>) {
+  // CHECK: %[[CAST:.+]] = builtin.unrealized_conversion_cast %arg0 : vector<i32> to vector<1xi32>
+  // CHECK: "llvm.intr.abs"(%[[CAST]]) <{is_int_min_poison = false}> : (vector<1xi32>) -> vector<1xi32>
+  %0 = math.absi %arg0 : vector<i32>
+  func.return
+}
+
+// -----
+
 // CHECK-LABEL: func @log1p(
 // CHECK-SAME: f32
 func.func @log1p(%arg0 : f32) {
@@ -201,6 +214,15 @@ func.func @ctlz(%arg0 : i32) {
   func.return
 }
 
+// CHECK-LABEL: func @ctlz_0d_vec(
+// CHECK-SAME: i32
+func.func @ctlz_0d_vec(%arg0 : vector<i32>) {
+  // CHECK: %[[CAST:.+]] = builtin.unrealized_conversion_cast %arg0 : vector<i32> to vector<1xi32>
+  // CHECK: "llvm.intr.ctlz"(%[[CAST]]) <{is_zero_poison = false}> : (vector<1xi32>) -> vector<1xi32>
+  %0 = math.ctlz %arg0 : vector<i32>
+  func.return
+}
+
 // -----
 
 // CHECK-LABEL: func @cttz(
@@ -213,6 +235,17 @@ func.func @cttz(%arg0 : i32) {
 
 // -----
 
+// CHECK-LABEL: func @cttz_0d_vec(
+// CHECK-SAME: i32
+func.func @cttz_0d_vec(%arg0 : vector<i32>) {
+  // CHECK: %[[CAST:.+]] = builtin.unrealized_conversion_cast %arg0 : vector<i32> to vector<1xi32>
+  // CHECK: "llvm.intr.cttz"(%[[CAST]]) <{is_zero_poison = false}> : (vector<1xi32>) -> vector<1xi32>
+  %0 = math.cttz %arg0 : vector<i32>
+  func.return
+}
+
+// -----
+
 // CHECK-LABEL: func @cttz_vec(
 // CHECK-SAME: i32
 func.func @cttz_vec(%arg0 : vector<4xi32>) {

[AGindinson]

@banach-space, @dcaballe, @Groverkss, @vzakhari, could you please take a look?

[banach-space]

Hi Artem - thanks for sending this!

Rank-0 vectors are tricky. Support across MLIR is still somewhat patchy, and our recent discussion on their role in the broader ecosystem was inconclusive:

• https://discourse.llvm.org/t/rfc-should-we-restrict-the-usage-of-0-d-vectors-in-the-vector-dialect

I’m bringing this up because I don’t see existing examples of rank-0 vectors in math-to-llvm.mlir, which suggests this may not have been explored yet in the context of MathToLLVM.

A couple of high-level questions:

• In practice, is there anything that will eliminate the builtin.unrealized_conversion_cast being inserted here? If not, could we instead lower the rank-0 vector to a scalar earlier in the pipeline?
• Why the focus on absi, ctlz, and cttz specifically? I’d expect this to apply to most (if not all) math ops that accept vector inputs.

In principle, I’m not opposed to this change - MLIR does support rank-0 vectors - but I want to make sure we consider this holistically. Ideally, we’d ensure that this integrates cleanly with the rest of the stack and doesn’t introduce hard-to-track edge cases later.

[AGindinson]

Thanks Andrzej for your detailed feedback!

Firstly, I absolutely should expand my proposal to other lowering patterns – this one in particular was the one showing up in my experiments with the IREE pipeline, but I'll see if I can trigger similar conversion errors through other E2E cases (or just in LIT mode).

There's clearly a lot of broader context around these scalar pseudo-vectors that I need to wrap my head around. My intuition has been that with mainstream dialect conversion, it should be possible to convert valid IR into valid IR – so until a hypothetical 0D-vec ban is imposed onto the Math specification, we should have at least some way to follow through with compilations.

Within an isolated pipeline, it would make perfect sense to me if Vector's in-house lowering of 0D to scalars took place as early as feasible, before MathToLLVM et al. Is there a way to guarantee that for all user pipelines, though? (Otherwise, the point about conformant IR -> conformant IR lowering requirements comes up again.)

I'm very keen to resolve the conversion issues in the cleanest possible way, so I'd really appreciate further guidance here. As an MLIR newcomer I haven't grown strong opinions on philosophical Discourses just yet :) So ready to implement anything that improves conformance across dialects rapidly enough.

[banach-space]

it should be possible to convert valid IR into valid IR

Agreed, in principle. My question is more about how to enable this. Should we insert:

• builtin.unrealized_conversion_cast %arg0 : vector<i32> to vector<1xi32> (i.e., "view" rank-0 vectors as rank-1)?
• builtin.unrealized_conversion_cast %arg0 : vector<i32> to i32 (i.e., treat rank-0 vectors as scalars)?
• Or insert vector.extract f32 from vector<f32> before hitting ConvertMathToLLVM? (That would require a dependency on Vector.)

My concern is that builtin.unrealized_conversion_cast can allow us to defer important design discussions. For example, I’d argue both vector<i32> and vector<1xi32> are sub-optimal if we're really just handling a scalar - in that case, shouldn't we lower it as a scalar?

That said, I realise you’re following the precedent set by the LLVM type converter:

llvm-project/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

Line 630 in 926c201

/// * 0-D `vector<T>` are converted to vector<1xT>

So let me reframe the question: is converting vector<f32> --> vector<1xf32> the right choice in the context of the Math dialect?

Otherwise, the point about conformant IR -> conformant IR lowering requirements comes up again.

Sure, though I think there's an implicit assumption here: that any form of Math or Vector should lower cleanly to LLVM. I don’t fully agree - and rank-0 vectors are a good example. LLVM doesn’t support them directly, so we need some decomposition or transformation (just as we do for vector.contract despite having ConvertVectorToLLVM).

To be clear: I’m not blocking the current patch 🙂 If builtin.unrealized_conversion_cast works - and often it does, since it cancels out with other casts - that’s fine by me.

I'd really appreciate further guidance here.

If you don’t have strong opinions (or time to dig into all these questions 😅), I’d say proceed with what you have. My only asks would be:

• Leave TODOs for other unhandled ops (bonus points if you tackle them).
• Add a comment in MathToLLVM.cpp to explain the need for this conversion (e.g., "required to legalize rank-0 vectors").

Lastly, just to double-check - does builtin.unrealized_conversion_cast work in your E2E flow?

Thanks again!

[AGindinson]

@banach-space All very good points, thanks. I'll go forth with implementing and/or clarifying the TODOs.

Just to double-check - does builtin.unrealized_conversion_cast work in your E2E flow?

It does, luckily :)

I realise you’re following the precedent set by the LLVM type converter

That's a neat way to decompose this, because my current approach is essentially deferring the work to the LLVM TypeConverter. And when we turn to the general matter, do any dialects/pipelines have an argument for maintaining the scalars in 1-d vector form at the LLVM level? Because if not, we could next implement one of the following depending on the design discussions:

• either propose a change in the core converter (vector<i32> to i32)
• or make sure there's a general, re-usable pass for "inserting vector.extract f32 from vector<f32>"

This wouldn't necessarily alleviate the need to explicitly handle 0-d vectors in each individual dialect (might even break some "manual" lowering patterns - need to check & see if the fixes are expensive), but that could hardly be a regression from the current state.

(*) E.g. even if some wanted to convert the LLVM dialect back into a higher-level one that would somehow really benefit from 0d vectors, they'd need to infer these semantics from vector<1xT> anyhow by looking at memory effects? At the end of the day, the pseudo-vectors will get conflated with either scalars or actual 1-elt vectors.

[Groverkss]

it should be possible to convert valid IR into valid IR

Agreed, in principle. My question is more about how to enable this. Should we insert:

• builtin.unrealized_conversion_cast %arg0 : vector<i32> to vector<1xi32> (i.e., "view" rank-0 vectors as rank-1)?
• builtin.unrealized_conversion_cast %arg0 : vector<i32> to i32 (i.e., treat rank-0 vectors as scalars)?
• Or insert vector.extract f32 from vector<f32> before hitting ConvertMathToLLVM? (That would require a dependency on Vector.)

My concern is that builtin.unrealized_conversion_cast can allow us to defer important design discussions. For example, I’d argue both vector<i32> and vector<1xi32> are sub-optimal if we're really just handling a scalar - in that case, shouldn't we lower it as a scalar?

That said, I realise you’re following the precedent set by the LLVM type converter:

llvm-project/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

Line 630 in 926c201

/// * 0-D `vector<T>` are converted to vector<1xT>

So let me reframe the question: is converting vector<f32> --> vector<1xf32> the right choice in the context of the Math dialect?

Otherwise, the point about conformant IR -> conformant IR lowering requirements comes up again.

Sure, though I think there's an implicit assumption here: that any form of Math or Vector should lower cleanly to LLVM. I don’t fully agree - and rank-0 vectors are a good example. LLVM doesn’t support them directly, so we need some decomposition or transformation (just as we do for vector.contract despite having ConvertVectorToLLVM).

To be clear: I’m not blocking the current patch 🙂 If builtin.unrealized_conversion_cast works - and often it does, since it cancels out with other casts - that’s fine by me.

I'd really appreciate further guidance here.

If you don’t have strong opinions (or time to dig into all these questions 😅), I’d say proceed with what you have. My only asks would be:

• Leave TODOs for other unhandled ops (bonus points if you tackle them).
• Add a comment in MathToLLVM.cpp to explain the need for this conversion (e.g., "required to legalize rank-0 vectors").

Lastly, just to double-check - does builtin.unrealized_conversion_cast work in your E2E flow?

Thanks again!

The current way of handeling this is correct and consistent. The LLVM type converter explicitly documents how vector types are converted to LLVM: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp#L629 . If handled consistently, the unrealized_conversion_casts will always cancel out.

I don't think we should bring the topic of what is the correct way here, we can bring it in the 0-D discourse discussion. This PR is consistent with the rest of the lowerings.

[AGindinson]

@banach-space, I've gone ahead and added the type conversions for the remaining "boilerplate" intrinsics.
@Groverkss, I've applied your comment and also made an effort to improve the consistency further. Please take a look if that works.

[Groverkss]

LGTM, please wait for @banach-space to have a look as well

[banach-space]

I am OK with these changes, but still don't understand the selection process for the Ops to update:

Employ this rule within the conversion pattern for: math.ctlz .cttz, .absi, .expm1, .log1p, .rsqrt, .isnan, .isfinite.

What about the other Ops? I am away next week, so please feel free to land this once my final comment is addressed, thanks!

[AGindinson]

I am OK with these changes, but still don't understand the selection process for the Ops to update:

Employ this rule within the conversion pattern for: math.ctlz .cttz, .absi, .expm1, .log1p, .rsqrt, .isnan, .isfinite.

What about the other Ops?

@banach-space Basically I've gone for all ops that are explicitly converted within MathToLLVM.cpp – apparently, the rest is fully covered by inheriting a bunch of patterns from Vector.

Thanks for the review Andrzej, Kunwar! If/when this is good to go in, I'd really appreciate help with the merge, as I haven't achieved write access just yet.

[banach-space]

@banach-space Basically I've gone for all ops that are explicitly converted within MathToLLVM.cpp – apparently, the rest is fully covered by inheriting a bunch of patterns from Vector.

Thanks! Please include this rationale in the summary - that's valuable context. Also, are those other conversions tested for rank-0? Just curious, not expecting you do add tests for those.

[AGindinson]

Please include this rationale in the summary - that's valuable context. Also, are those other conversions tested for rank-0? Just curious, not expecting you do add tests for those.

Updated the description (also correcting myself along the way – these are from LLVMCommon/VectorPattern). I haven't been able to find any tests at all for the inherited patterns within Math (but I'd need a second look). That would be somewhat of a miss...

[github-actions]

@AGindinson Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!