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

[nicolasvasilache]

Summary

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

[nicolasvasilache]

Rebase.

[gysit]

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

1161

nit: should

1167

nit: auto format wraps the lines in at an unfortunate position. maybe everything fits one line after removing the indention.

[nicolasvasilache]

Address.

[nicolasvasilache]

Fix test.

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

static ParseResult parseExtractValueOp(OpAsmParser &parser,
   return success();
 }
 
 OpFoldResult LLVM::ExtractValueOp::fold(ArrayRef<Attribute> operands) {
   auto insertValueOp = getContainer().getDefiningOp<InsertValueOp>();
   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 {};
 }
 
 static LogicalResult verify(ExtractValueOp op) {
   Type valueType = getInsertExtractValueElementType(op.getContainer().getType(),
                                                     op.getPositionAttr(), op);

mlir/test/Dialect/LLVMIR/canonicalize.mlir

llvm.func @fold_extractvalue() -> i32 {
   // CHECK-NOT: extractvalue
   %3 = llvm.extractvalue %2[0] : !llvm.struct<(i32, i32)>
   %4 = llvm.extractvalue %2[1] : !llvm.struct<(i32, i32)>
 
   //     CHECK: llvm.add %[[C0]], %[[C1]]
   %5 = llvm.add %3, %4 : i32
   llvm.return %5 : i32
 }
+
+// -----
+
+// CHECK-LABEL: no_fold_extractvalue
+llvm.func @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
+}