[Mlir-commits] [mlir] 4a94a33 - [MLIR][LLVM] Fold extractvalue to ignore insertvalue at distinct index

[William S. Moses]

Author: William S. Moses
Date: 2022-03-04T11:03:34-05:00
New Revision: 4a94a33ca6cb9098ed665dac74f418615d35ddea

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

LOG: [MLIR][LLVM] Fold extractvalue to ignore insertvalue at distinct index

We can simplify an extractvalue of an insertvalue to extract out of the base of the insertvalue, if the insert and extract are at distinct and non-prefix'd indices

Reviewed By: ftynse

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

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 3e855f5c01dc6..b3d3ab829dc85 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1341,6 +1341,7 @@ ParseResult ExtractValueOp::parse(OpAsmParser &parser, OperationState &result) {
 
 OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
   auto insertValueOp = getContainer().getDefiningOp<InsertValueOp>();
+  OpFoldResult result = {};
   while (insertValueOp) {
     if (getPosition() == insertValueOp.getPosition())
       return insertValueOp.getValue();
@@ -1358,10 +1359,16 @@ OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
     // ```
     if (getPosition().getValue().take_front(min) ==
         insertValueOp.getPosition().getValue().take_front(min))
-      return {};
+      return result;
+
+    // If neither a prefix, nor the exact position, we can extract out of the
+    // value being inserted into. Moreover, we can try again if that operand
+    // is itself an insertvalue expression.
+    getContainerMutable().assign(insertValueOp.getContainer());
+    result = getResult();
     insertValueOp = insertValueOp.getContainer().getDefiningOp<InsertValueOp>();
   }
-  return {};
+  return result;
 }
 
 LogicalResult ExtractValueOp::verify() {

diff  --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 5fbe0b648d48d..3e7b545875662 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -37,6 +37,18 @@ llvm.func @no_fold_extractvalue(%arr: !llvm.array<4xf32>) -> f32 {
   %3 = llvm.extractvalue %2[0, 0] : !llvm.array<4 x !llvm.array<4xf32>>
 
   llvm.return %3 : f32
+
+}
+// -----
+
+// CHECK-LABEL: fold_unrelated_extractvalue
+llvm.func @fold_unrelated_extractvalue(%arr: !llvm.array<4xf32>) -> f32 {
+  %f0 = arith.constant 0.0 : f32
+  // CHECK-NOT: insertvalue
+  // CHECK: extractvalue
+  %2 = llvm.insertvalue %f0, %arr[0] : !llvm.array<4xf32>
+  %3 = llvm.extractvalue %2[1] : !llvm.array<4xf32>
+  llvm.return %3 : f32
 }
 
 // -----