[Mlir-commits] [mlir] [MLIR][LLVM] Fold extract of constant (PR #127927)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: William Moses (wsmoses)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/127927.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp (+13) 
- (modified) mlir/test/Dialect/LLVMIR/canonicalize.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index bfcba40555a7c..23fbe38467b6a 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1932,6 +1932,19 @@ OpFoldResult LLVM::ExtractValueOp::fold(FoldAdaptor adaptor) {
     getContainerMutable().set(extractValueOp.getContainer());
     return getResult();
   }
+
+  {
+    DenseElementsAttr constval;
+    matchPattern(getContainer().getOperand(), m_Constant(&constval));
+    if (constval && constval.getElementType() == getType()) {
+      if (constval.isSplat()) {
+        return constval.getSplatValue();
+      } else if (getPosition().size() == 1) {
+        return constval.getValues()[getPosition()[0]];
+      }
+    }
+  }
+
   auto insertValueOp = getContainer().getDefiningOp<InsertValueOp>();
   OpFoldResult result = {};
   while (insertValueOp) {
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index c509cd82227c2..055cb38c65337 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -94,6 +94,26 @@ llvm.func @fold_extract_extractvalue(%arr: !llvm.struct<(i64, array<1 x ptr<1>>)
   llvm.return %b : !llvm.ptr<1>
 }
 
+// -----
+// CHECK-LABEL: fold_extract_const
+// CHECK-NOT: extractvalue
+// CHECK: llvm.mlir.constant(-5.0)
+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.9)
+llvm.func @fold_extract_const() -> 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

``````````

</details>


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