[Mlir-commits] [mlir] 7a67e41 - [mlir][LLVM] Fix incorrect GEP fold with struct constants

[Markus Böck]

Author: Markus Böck
Date: 2022-07-27T14:42:35+02:00
New Revision: 7a67e41be2a56df8629cdb3f6b529855ecce3ed4

URL: https://github.com/llvm/llvm-project/commit/7a67e41be2a56df8629cdb3f6b529855ecce3ed4
DIFF: https://github.com/llvm/llvm-project/commit/7a67e41be2a56df8629cdb3f6b529855ecce3ed4.diff

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

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.

Differential Revision: https://reviews.llvm.org/D129374

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
    mlir/test/Dialect/LLVMIR/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index b3f8012c496cd..f635c9c214342 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2699,7 +2699,7 @@ OpFoldResult LLVM::AddrSpaceCastOp::fold(ArrayRef<Attribute> operands) {
 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 {};
 }

diff  --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 8557a51a2a888..459fa2f8d3627 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -100,6 +100,17 @@ llvm.func @fold_gep(%x : !llvm.ptr<i8>) -> !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