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

[Tobias Gysi]

================
@@ -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.
+    if (!attrElement)
+      return containerAttr;
----------------
gysit wrote:

Would it make sense to check the type of the returned attribute here to ensure we only return ZeroAttr, UndefAttr etc but not an ArrayAttr for example? 


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