[Mlir-commits] [mlir] [MLIR] Extend the extractvalue fold method (PR #172297)

[Vadim Curcă]

================
@@ -1975,9 +1981,32 @@ OpFoldResult LLVM::ExtractValueOp::fold(FoldAdaptor adaptor) {
       getContainerMutable().assign(insertValueOp.getContainer());
       result = getResult();
     }
-    insertValueOp = insertValueOp.getContainer().getDefiningOp<InsertValueOp>();
+    container = insertValueOp.getContainer().getDefiningOp();
+  }
+  if (!container)
+    return result;
+
+  Attribute containerAttr;
+  if (!matchPattern(container, m_Constant(&containerAttr)))
+    return nullptr;
+  for (int64_t pos : extractPos) {
+    Attribute attrElement = extractElementAt(containerAttr, pos);
+
+    // It is possible to fail to extract an element from the container and still
+    // fold the operation to a constant. For example:
+    // ```
+    // %container = llvm.mlir.zero : !llvm.struct<(i8, i32)>
+    // %result = llvm.extractvalue %container[0] : !llvm.struct<(i8, i32)>
+    // ```
+    // In this case, `containerAttr` is an `LLVM::ZeroAttr` that does not
+    // contain any nested elements, yet the operation can be folded to a zero
+    // constant.
----------------
VadimCurca wrote:

I didn't know how to check that `containerAttr` is a valid attribute in a clean way.

The folder can return an attribute that cannot be used to fold the operation. The greedy pattern rewriter driver will try to materialize a constant using the returned attribute and the result type of the folded operation ([GreedyPatternRewriteDriver.cpp#L530](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp#L530)); if it fails, the operation is not replaced.

Do you have suggestions for how to do this?

https://github.com/llvm/llvm-project/pull/172297