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

[William Moses]

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

>From 1daf6c612ae3defd264c89aa745714e453717393 Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh at wsmoses.com>
Date: Wed, 19 Feb 2025 20:04:17 -0500
Subject: [PATCH 1/2] [MLIR][LLVM] Fold extract of constant

---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 14 ++++++++++++++
 mlir/test/Dialect/LLVMIR/canonicalize.mlir | 20 ++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index bfcba40555a7c..131d8070717f3 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1932,6 +1932,20 @@ 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) {
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index c509cd82227c2..6b7bec3d23dff 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.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

>From 3d64d7d37131455a82731019257d5c1061c4200c Mon Sep 17 00:00:00 2001
From: William Moses <gh at wsmoses.com>
Date: Thu, 20 Feb 2025 10:00:22 -0600
Subject: [PATCH 2/2] Apply suggestions from code review

Co-authored-by: Christian Ulmann <christianulmann at gmail.com>
---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp | 6 ++----
 mlir/test/Dialect/LLVMIR/canonicalize.mlir | 2 ++
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 131d8070717f3..ccf8f72b2b230 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1937,12 +1937,10 @@ OpFoldResult LLVM::ExtractValueOp::fold(FoldAdaptor adaptor) {
     DenseElementsAttr constval;
     matchPattern(getContainer(), m_Constant(&constval));
     if (constval && constval.getElementType() == getType()) {
-      if (isa<SplatElementsAttr>(constval)) {
+      if (isa<SplatElementsAttr>(constval))
         return constval.getSplatValue<Attribute>();
-      }
-      if (getPosition().size() == 1) {
+      if (getPosition().size() == 1)
         return constval.getValues<Attribute>()[getPosition()[0]];
-      }
     }
   }
 
diff --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 6b7bec3d23dff..a793caca064ec 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -95,6 +95,7 @@ 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)
@@ -105,6 +106,7 @@ llvm.func @fold_extract_const() -> f64 {
 }
 
 // -----
+
 // CHECK-LABEL: fold_extract_splat
 // CHECK-NOT: extractvalue
 // CHECK: llvm.mlir.constant(-8.900000e+01 : f64)