[mlir][llvm] Fix elem type passing into getelementptr

[rikhuijzer]

As was correctly pointed out by @azteca1998, the element type for a llvm.getelementptr was only read when using an attribute and not when using a type. As pointed out in #63832 (comment), the translation to LLVM would work for

llvm.func @main(%0 : !llvm.ptr) -> !llvm.ptr {
    %1 = llvm.getelementptr %0[0] { elem_type = !llvm.ptr } : (!llvm.ptr) -> !llvm.ptr
    llvm.return %1 : !llvm.ptr
}

but not for

llvm.func @main(%0 : !llvm.ptr) -> !llvm.ptr<ptr> {
    %1 = llvm.getelementptr %0[0] : (!llvm.ptr) -> !llvm.ptr<ptr>
    llvm.return %1 : !llvm.ptr<ptr>
}

This was caused by the LLVM_GEPOp builder only reading the type from the attribute ({ elem_type = !llvm.ptr }), but not from the pointer type (!llvm.ptr<ptr>).

Fixes #63832.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Changes

As was correctly pointed out by @azteca1998, the element type for a llvm.getelementptr was only correctly read when using an attribute and not when using a type. As pointed out in #63832 (comment), the translation to LLVM would work for

llvm.func @<!-- -->main(%0 : !llvm.ptr) -&gt; !llvm.ptr {
    %1 = llvm.getelementptr %0[0] { elem_type = !llvm.ptr } : (!llvm.ptr) -&gt; !llvm.ptr
    llvm.return %1 : !llvm.ptr
}

but not for

llvm.func @<!-- -->main(%0 : !llvm.ptr) -&gt; !llvm.ptr&lt;ptr&gt; {
    %1 = llvm.getelementptr %0[0] : (!llvm.ptr) -&gt; !llvm.ptr&lt;ptr&gt;
    llvm.return %1 : !llvm.ptr&lt;ptr&gt;
}

This was caused by the LLVM_GEPOp builder only reading the type from the attribute ({ elem_type = !llvm.ptr }), but not from the pointer type (!llvm.ptr<ptr>).

Fixes #63832.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td (+7-2)
• (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp (+6-3)
• (modified) mlir/test/Target/LLVMIR/opaque-ptr.mlir (+9)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 25209ce4497455e..9d46c5b3bdc5e1b 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -303,8 +303,13 @@ def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
         indices.push_back(
             builder.getInt32(valueOrAttr.get<IntegerAttr>().getInt()));
     }
-    Type baseElementType = op.getSourceElementType();
-    llvm::Type *elementType = moduleTranslation.convertType(baseElementType);
+
+    Type elemTypeFromAttr = op.getSourceElementType();
+    auto ptrType = ::llvm::cast<LLVMPointerType>(op.getType());
+    Type elemTypeFromPtrType = ptrType.getElementType();
+
+    llvm::Type *elementType = moduleTranslation.convertType(
+        elemTypeFromAttr ? elemTypeFromAttr : elemTypeFromPtrType);
     $res = builder.CreateGEP(elementType, $base, indices, "", $inbounds);
   }];
   let assemblyFormat = [{
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 95c04098d05fc2f..62cb595069e6652 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -287,14 +287,17 @@ ParseResult AllocaOp::parse(OpAsmParser &parser, OperationState &result) {
 }
 
 /// Checks that the elemental type is present in either the pointer type or
-/// the attribute, but not both.
+/// the attribute, but not in none or both.
 static LogicalResult verifyOpaquePtr(Operation *op, LLVMPointerType ptrType,
                                      std::optional<Type> ptrElementType) {
-  if (ptrType.isOpaque() && !ptrElementType.has_value()) {
+  bool typePresentInPointerType = !ptrType.isOpaque();
+  bool typePresentInAttribute = ptrElementType.has_value();
+
+  if (!typePresentInPointerType && !typePresentInAttribute) {
     return op->emitOpError() << "expected '" << kElemTypeAttrName
                              << "' attribute if opaque pointer type is used";
   }
-  if (!ptrType.isOpaque() && ptrElementType.has_value()) {
+  if (typePresentInPointerType && typePresentInAttribute) {
     return op->emitOpError()
            << "unexpected '" << kElemTypeAttrName
            << "' attribute when non-opaque pointer type is used";
diff --git a/mlir/test/Target/LLVMIR/opaque-ptr.mlir b/mlir/test/Target/LLVMIR/opaque-ptr.mlir
index c21f9b0542debc6..3bde192b4cc4d02 100644
--- a/mlir/test/Target/LLVMIR/opaque-ptr.mlir
+++ b/mlir/test/Target/LLVMIR/opaque-ptr.mlir
@@ -42,6 +42,15 @@ llvm.func @opaque_ptr_gep_struct(%arg0: !llvm.ptr, %arg1: i32) -> !llvm.ptr {
   llvm.return %0 : !llvm.ptr
 }
 
+// CHECK-LABEL: @opaque_ptr_elem_type
+llvm.func @opaque_ptr_elem_type(%0: !llvm.ptr) -> !llvm.ptr {
+  // CHECK: getelementptr ptr, ptr
+  %1 = llvm.getelementptr %0[0] { elem_type = !llvm.ptr } : (!llvm.ptr) -> !llvm.ptr
+  // CHECK: getelementptr ptr, ptr
+  %2 = llvm.getelementptr %0[0] : (!llvm.ptr) -> !llvm.ptr<ptr>
+  llvm.return %1 : !llvm.ptr
+}
+
 // CHECK-LABEL: @opaque_ptr_matrix_load_store
 llvm.func @opaque_ptr_matrix_load_store(%ptr: !llvm.ptr, %stride: i64) -> vector<48 x f32> {
   // CHECK: call <48 x float> @llvm.matrix.column.major.load.v48f32.i64

[zero9178]

What are the use cases for allowing either of these two syntaxes?
The LLVM Dialect tries to closely mirror LLVM proper as much as possible and this would deviate from LLVMs behaviour. While the transition is currently stalled, in the future typed pointers will be removed from the dialect entirely, further discouraging the use of typed-pointers such as shown here.

[rikhuijzer]

What are the use cases for allowing either of these two syntaxes?
The LLVM Dialect tries to closely mirror LLVM proper as much as possible and this would deviate from LLVMs behaviour. While the transition is currently stalled, in the future typed pointers will be removed from the dialect entirely, further discouraging the use of typed-pointers such as shown here.

From reading https://reviews.llvm.org/D123310, I assume that the use-case is that "MLIR can afford to support both opaque and non-opaque pointers at the same time and simplify the transition". But Alex (@ftynse) is probably the best person to answer this.

Regardless of the big picture, the question might be unrelated to this PR as this PR only fixes a bug in the implementation. Currently, there are LLVM operations in MLIR that do accept both forms. For example, llvm.alloca where

llvm.func @main(%sz : i64) {
  "llvm.alloca"(%sz) : (i64) -> !llvm.ptr<i32>
  llvm.return
}

and

llvm.func @main(%sz : i64) {
  "llvm.alloca"(%sz) { elem_type = i32 } : (i64) -> !llvm.ptr
  llvm.return
}

Both are translated to

; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"

declare ptr @malloc(i64)

declare void @free(ptr)

define void @main(i64 %0) {
  %2 = alloca i32, i64 %0, align 4
  ret void
}

!llvm.module.flags = !{!0}

!0 = !{i32 2, !"Debug Info Version", i32 3}

via

mlir-translate temp.mlir -mlir-to-llvmir

Conversely, llvm.getelementptr currently only accepts the attribute ({ elem_type = !llvm.ptr }). This PR makes things consistent again.

[zero9178]

Thank you for your elaboration!
I am not sure whether the original commit is referring to also supporting both typed pointers and opaque pointers in the exact same IR. This behaviour deviates from LLVM entirely as LLVM only supported one or the other as well.

Furthermore, your example of the alloca is not quite equivalent to the GEP. Your first example uses strictly typed pointers while the second uses strictly opaque pointers. The GEP llvm.getelementptr %0[0] : (!llvm.ptr) -> !llvm.ptr<ptr> mixes both opaque pointers and typed pointers, which has so far not been supported.
Note that: llvm.getelementptr %0[0] : (!llvm.ptr) -> !llvm.ptr, !llvm.ptr and llvm.getelementptr %0[0] : (!llvm.ptr<ptr>) -> !llvm.ptr<ptr> work without issues.

Lastly, the PR also deviates from how GEPs works with typed pointers as it deduces the element type from the result pointer. If we really want to support a half-typed half-opaque GEP it'd make more sense to support llvm.getelementptr %0[0] : (!llvm.ptr<ptr>) -> !llvm.ptr. This would also allow specifying more than one index, something the current solution doesn't. E.g.

`llvm.getelementptr %0[0, 1] : (!llvm.ptr<struct<(i32, i8)>>) -> !llvm.ptr`

With all that said, I see this more as a "potential missing feature" rather than a bug. That is why I am curious what the motivation here is. Is this to ease transition to opaque pointers? I am not sure how effective this would be and I am a bit afraid of the burden of having to maintain what is essentially a third GEP representation.

[rikhuijzer]

[...] I am a bit afraid of the burden of having to maintain what is essentially a third GEP representation.

Thanks for the explanation @zero9178. I am affraid I do not fully understand your reasoning (most likely due to a lack of knowledge on my side). However, it sounds very reasonable to me!

So before I spend more time on the code, shall I rewrite this PR to throw a clear error from GEPOp::verify for the incorrect representation? Then we can close #63832.

[zero9178]

So before I spend more time on the code, shall I rewrite this PR to throw a clear error from GEPOp::verify for the incorrect representation? Then we can close #63832.

I'd personally be in favour of that. There is a lot of room for improvements in the verifiers in the LLVM Dialect so any contributions there are definitely highly appreciated!
That said, its probably best to wait a bit for others to comment as well as they may have different views than I do.

[ftynse]

What I meant in the comment was that MLIR doesn't need to make a hard switch to using opaque pointers and immediately turn off typed pointers. I didn't really imply that an operations in the IR should freely mix the two. Typed pointers should and will be removed from the LLVM dialect as well. I consider any significant investment into typed pointer support, let alone a "mixed mode", a waste of time. Since this PR fixes a crash, I'm fine with it landing as is instead of spending cycles on an arguably better fix that will get removed within months as I hope opaque pointers to be gone from the dialect by the next release. Improving verifiers otherwise is greatly appreciated!

[rikhuijzer]

Thanks both for the review and clarifications. I'm learning a lot.

[aartbik]

This broke the bot?

https://lab.llvm.org/buildbot/#/builders/61/builds/50100

[joker-eph]

Reverted, the bots have been broken all day.

[stellaraccident]

Thank you for the revert.

We were failing on assertion when integrating this patch:

iree-compile: /home/stella/megabump/work/iree/third_party/llvm-project/llvm/include/llvm/Support/Casting.h:566: decltype(auto) llvm::cast(const From &) [To = mlir::LLVM::LLVMPointerType, F
rom = mlir::Type]: Assertion `isa<To>(Val) && "cast<Ty>() argument of incompatible type!"' failed.