[mlir][LLVM] Fix incorrect GEP fold with struct constants

[zero9178]

Summary

The fold in it's current state only checks whether the amount of dynamic indices is 1. This does however not check for the presence of any struct indices, leading to an incorrect fold.

This patch fixes that issue by checking that struct indices are 1, which in addition to the pre-existing check that dynamic indices are 1, guarantees that the single index is a dynamic one.

[zero9178]

Friendly ping :)

[ftynse]

I can't match the test to the code, or to the description. I would expect the first condition (getBase.getType() == getType()) to fail here because the type of base is ptr-to-struct, and the type of the result is ptr-to-i32. (Note that something may be broken because of the opaque pointer transition, in which case both types are opaque pointers, and this needs a separate fix). Also, looking at the comment, it looks like the condition is to have the _total_ number of indices (i.e. the sum of sizes of indices and struct indices) should be one, the code checks for both... Something strange is going on here.

[zero9178]

• Use opaque pointers in test, to properly test the problematic part of the if condition
• Just check struct indices, not both

[zero9178]

In D129374#3665376, @ftynse wrote:

I can't match the test to the code, or to the description. I would expect the first condition (getBase.getType() == getType()) to fail here because the type of base is ptr-to-struct, and the type of the result is ptr-to-i32. (Note that something may be broken because of the opaque pointer transition, in which case both types are opaque pointers, and this needs a separate fix). Also, looking at the comment, it looks like the condition is to have the _total_ number of indices (i.e. the sum of sizes of indices and struct indices) should be one, the code checks for both... Something strange is going on here.

I am not sure how it originally failed, but I changed it to using opaque pointers now, in which case it's obvious that the condition is true (and also matches the IR I originally found the bug in downstream).
The reason I initially checked for both indices is because I mistakenly thought that structIndices could be 1, but indices 0, but that is not true since the first index is never a struct index.
I changed it to just use the size of struct indices now since that is always equal to the total amount of indices (meanwhile indices is only ever equal to the amount of dynamic ones)

[ftynse]

This still looks incorrect. The condition may now segfault when getIndices().size() == 0. I think you really want to check getStructIndices().empty() and getIndices().size() == 1 because that's where this folding applies: GEP of a pointer with zero offset folds to the same pointer. Any more complicated GEP shouldn't fold.

[zero9178]

In D129374#3667671, @ftynse wrote:

This still looks incorrect. The condition may now segfault when getIndices().size() == 0. I think you really want to check getStructIndices().empty() and getIndices().size() == 1 because that's where this folding applies: GEP of a pointer with zero offset folds to the same pointer. Any more complicated GEP shouldn't fold.

I was under impression that struct indices is always equal to the total amount of indices (dynamic + constants for structs) in a GEP instruction, and that struct indices that are actually dynamic are simply set to kDynamicIndex.
This is even checked in an assert here: https://github.com/llvm/llvm-project/blob/a2158374ba1a6f81f4cce3eb54d0bc44f3ab75e0/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp#L594
Hence the check for structIndices to be sized 1 is required. If it we were to check whether it is empty it leads to the fold never working and the @fold_gep test above failing.
Additionally a check for getIndices().size() == 1 is also theoretically unnecessary as the very first index is always dynamic. This is already assumed elsewhere in the verifier code eg. here:
https://github.com/llvm/llvm-project/blob/a2158374ba1a6f81f4cce3eb54d0bc44f3ab75e0/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp#L546
(it should also be verified and not just assumed, but that is another bug and issue and not relevant for this patch).
Nevertheless I'd we can leave the getIndices().size() == 1 check in regardless unless you prefer otherwise.

[zero9178]

Check that both struct indices and normal indices are of size 1

[zero9178]

In D129374#3669868, @zero9178 wrote:

...
Additionally a check for getIndices().size() == 1 is also theoretically unnecessary as the very first index is always dynamic. This is already assumed elsewhere in the verifier code eg. here:
https://github.com/llvm/llvm-project/blob/a2158374ba1a6f81f4cce3eb54d0bc44f3ab75e0/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp#L546
(it should also be verified and not just assumed, but that is another bug and issue and not relevant for this patch).
Nevertheless I'd we can leave the getIndices().size() == 1 check in regardless unless you prefer otherwise.

Actually this part is additionally more confusing since it works if you parse llvm.getelementptr %c[0] eg. but using it's build methods it is not possible to create such a GEPOp as it will hit that faulty code doing an out of bounds index. So there is some sort of bug here (either in the verifier, or the builder) and the correct fix will depend on whether using constant indices for things that don't require them should be allowed or not.

Regardless of that though, the patch should be correct now in the current state (and whatever the outcome of the above may be)

[ftynse]

There's some tech debt here... I don't have a strong opinion on whether to allow attribute indices for non-struct. We may allow them as long as it does not complexity the verification and translation. We may also want a builder/getter interface that uses OpFoldResult (yes, the naming is annoying) that is actually a union of Attribute and Value as a way to provide a single unified list of indices and stop exposing kDynamicIndex and friends.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::GEPOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult LLVM::GEPOp::fold(ArrayRef<Attribute> operands) {
   // gep %x:T, 0 -> %x
   if (getBase().getType() == getType() && getIndices().size() == 1 &&
-      matchPattern(getIndices()[0], m_Zero()))
+      getStructIndices().size() == 1 && matchPattern(getIndices()[0], m_Zero()))
     return getBase();
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/LLVMIR/canonicalize.mlir

 // CHECK-SAME: %[[a0:arg[0-9]+]]
 // CHECK-NEXT: llvm.return %[[a0]]
 llvm.func @fold_gep(%x : !llvm.ptr<i8>) -> !llvm.ptr<i8> {
   %c0 = arith.constant 0 : i32
   %c = llvm.getelementptr %x[%c0] : (!llvm.ptr<i8>, i32) -> !llvm.ptr<i8>
   llvm.return %c : !llvm.ptr<i8>
 }
 
+// CHECK-LABEL: fold_gep_neg
+// CHECK-SAME: %[[a0:arg[0-9]+]]
+// CHECK-NEXT: %[[C:.*]] = arith.constant 0
+// CHECK-NEXT: %[[RES:.*]] = llvm.getelementptr %[[a0]][%[[C]], 1]
+// CHECK-NEXT: llvm.return %[[RES]]
+llvm.func @fold_gep_neg(%x : !llvm.ptr) -> !llvm.ptr {
+  %c0 = arith.constant 0 : i32
+  %0 = llvm.getelementptr %x[%c0, 1] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.struct<(i32, i32)>
+  llvm.return %0 : !llvm.ptr
+}
+
 // -----
 
 // Check that LLVM constants participate in cross-dialect constant folding. The
 // resulting constant is created in the arith dialect because the last folded
 // operation belongs to it.
 // CHECK-LABEL: llvm_constant
 func.func @llvm_constant() -> i32 {
   // CHECK-NOT: llvm.mlir.constant