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

[wsmoses]

Summary

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

[wsmoses]

Fix rebase

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

if (!elementType)
     return failure();
 
   result.addTypes(elementType);
   return success();
 }
 
 OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
   auto insertValueOp = getContainer().getDefiningOp<InsertValueOp>();
+  OpFoldResult result = {};
   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 {};
+      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() {
   Type valueType = getInsertExtractValueElementType(getContainer().getType(),
                                                     getPositionAttr(), *this);
   if (!valueType)
     return failure();
 

mlir/test/Dialect/LLVMIR/canonicalize.mlir

llvm.func @no_fold_extractvalue(%arr: !llvm.array<4xf32>) -> f32 {
   // 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
+
+}
+// -----
+
+// 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
 }
 
 // -----
 
 // CHECK-LABEL: fold_bitcast
 // CHECK-SAME: %[[a0:arg[0-9]+]]
 // CHECK-NEXT: llvm.return %[[a0]]
 llvm.func @fold_bitcast(%x : !llvm.ptr<i8>) -> !llvm.ptr<i8> {