[Mlir-commits] [mlir] b377807 - [mlir][LLVM] Fix folding of LLVM::ExtractValueOp

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-11-16T14:49:05Z
New Revision: b377807a76e7ca9756d85f07907ba5e0b5b76e60

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

LOG: [mlir][LLVM] Fix folding of LLVM::ExtractValueOp

Limit the backtracking along def-use chains when a prefix is encountered as it would generate incorrect foldings.

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

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 ba7ec4f283af2..dd2ffdabd8c5f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1155,6 +1155,21 @@ OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
   while (insertValueOp) {
     if (getPosition() == insertValueOp.getPosition())
       return insertValueOp.getValue();
+    unsigned min =
+        std::min(getPosition().size(), insertValueOp.getPosition().size());
+    // If one is fully prefix of the other, stop propagating back as it will
+    // miss dependencies. For instance, %3 should not fold to %f0 in the
+    // following example:
+    // ```
+    //   %1 = llvm.insertvalue %f0, %0[0, 0] :
+    //     !llvm.array<4 x !llvm.array<4xf32>>
+    //   %2 = llvm.insertvalue %arr, %1[0] :
+    //     !llvm.array<4 x !llvm.array<4xf32>>
+    //   %3 = llvm.extractvalue %2[0, 0] : !llvm.array<4 x !llvm.array<4xf32>>
+    // ```
+    if (getPosition().getValue().take_front(min) ==
+        insertValueOp.getPosition().getValue().take_front(min))
+      return {};
     insertValueOp = insertValueOp.getContainer().getDefiningOp<InsertValueOp>();
   }
   return {};

diff  --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 7d21ff986c3d8..acbe0035e1470 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -21,3 +21,20 @@ llvm.func @fold_extractvalue() -> i32 {
   %5 = llvm.add %3, %4 : i32
   llvm.return %5 : i32
 }
+
+// -----
+
+// CHECK-LABEL: no_fold_extractvalue
+llvm.func @no_fold_extractvalue(%arr: !llvm.array<4xf32>) -> f32 {
+  %f0 = arith.constant 0.0 : f32
+  %0 = llvm.mlir.undef : !llvm.array<4 x !llvm.array<4xf32>>
+
+  // CHECK: insertvalue
+  // CHECK: insertvalue
+  // CHECK: extractvalue
+  %1 = llvm.insertvalue %f0, %0[0, 0] : !llvm.array<4 x !llvm.array<4xf32>>
+  %2 = llvm.insertvalue %arr, %1[0] : !llvm.array<4 x !llvm.array<4xf32>>
+  %3 = llvm.extractvalue %2[0, 0] : !llvm.array<4 x !llvm.array<4xf32>>
+
+  llvm.return %3 : f32
+}