[MLIR][LLVM] Fold extract of constant (#127927)

Co-authored-by: Christian Ulmann <[email protected]>

Author: wsmoses

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

@@ -1932,6 +1932,18 @@ OpFoldResult LLVM::ExtractValueOp::fold(FoldAdaptor adaptor) {
     getContainerMutable().set(extractValueOp.getContainer());
     return getResult();
   }
+
+  {
+    DenseElementsAttr constval;
+    matchPattern(getContainer(), m_Constant(&constval));
+    if (constval && constval.getElementType() == getType()) {
+      if (isa<SplatElementsAttr>(constval))
+        return constval.getSplatValue<Attribute>();
+      if (getPosition().size() == 1)
+        return constval.getValues<Attribute>()[getPosition()[0]];
+    }
+  }
+
   auto insertValueOp = getContainer().getDefiningOp<InsertValueOp>();
   OpFoldResult result = {};
   while (insertValueOp) {

mlir/test/Dialect/LLVMIR/canonicalize.mlir

@@ -96,6 +96,28 @@ llvm.func @fold_extract_extractvalue(%arr: !llvm.struct<(i64, array<1 x ptr<1>>)
 
 // -----
 
+// CHECK-LABEL: fold_extract_const
+// CHECK-NOT: extractvalue
+// CHECK:  llvm.mlir.constant(5.000000e-01 : f64)
+llvm.func @fold_extract_const() -> f64 {
+  %a = llvm.mlir.constant(dense<[-8.900000e+01, 5.000000e-01]> : tensor<2xf64>) : !llvm.array<2 x f64>
+  %b = llvm.extractvalue %a[1] : !llvm.array<2 x f64>
+  llvm.return %b : f64
+}
+
+// -----
+
+// CHECK-LABEL: fold_extract_splat
+// CHECK-NOT: extractvalue
+// CHECK: llvm.mlir.constant(-8.900000e+01 : f64)
+llvm.func @fold_extract_splat() -> f64 {
+  %a = llvm.mlir.constant(dense<-8.900000e+01> : tensor<2xf64>) : !llvm.array<2 x f64>
+  %b = llvm.extractvalue %a[1] : !llvm.array<2 x f64>
+  llvm.return %b : f64
+}
+
+// -----
+
 // CHECK-LABEL: fold_bitcast
 // CHECK-SAME: %[[ARG:[[:alnum:]]+]]
 // CHECK-NEXT: llvm.return %[[ARG]]