[Mlir-commits] [mlir] [mlir][tosa] Fold reshape of dense resource constants (PR #213225)

[Thibaut Goetghebuer-Planchon]

https://github.com/Tessil updated https://github.com/llvm/llvm-project/pull/213225

>From c51708ff066739a8f3921ff8a2e7e6dbd0cb8fd6 Mon Sep 17 00:00:00 2001
From: Thibaut Goetghebuer-Planchon <thibaut.goetghebuer-planchon at arm.com>
Date: Fri, 31 Jul 2026 09:29:33 +0100
Subject: [PATCH 1/2] [mlir][tosa] Fold reshape of dense resource constants

---
 .../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 10 ++++++++++
 mlir/test/Dialect/Tosa/constant_folding.mlir  | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index bb3c0f16b97aa..d3ded54f5d909 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2117,6 +2117,16 @@ OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
   if (!inputTy.getElementType().isIntOrIndexOrFloat())
     return {};
 
+  // Reshaping a resource-backed constant only requires updating its type.
+  if (auto operand = llvm::dyn_cast_if_present<DenseResourceElementsAttr>(
+          adaptor.getInput1())) {
+    if (!outputTy.hasStaticShape()) {
+      return {};
+    }
+
+    return DenseResourceElementsAttr::get(outputTy, operand.getRawHandle());
+  }
+
   // reshape(const(x)) -> const(reshape-attr(x))
   if (auto operand =
           llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1())) {
diff --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index 118746704d4c4..a2a0cab592220 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -889,6 +889,25 @@ func.func @reshape_splat() -> tensor<6x5x4xi32> {
 
 // -----
 
+// CHECK-LABEL: @reshape_dense_resource
+func.func @reshape_dense_resource() -> tensor<4xf32> {
+  // CHECK: %[[RESHAPED:.+]] = "tosa.const"() <{values = dense_resource<reshape_resource> : tensor<4xf32>}> : () -> tensor<4xf32>
+  %input = "tosa.const"() <{values = dense_resource<reshape_resource> : tensor<2x2xf32>}> : () -> tensor<2x2xf32>
+  %shape = tosa.const_shape {values = dense<4> : tensor<1xindex>} : () -> !tosa.shape<1>
+  %reshape = tosa.reshape %input, %shape : (tensor<2x2xf32>, !tosa.shape<1>) -> tensor<4xf32>
+  // CHECK: return %[[RESHAPED]]
+  return %reshape : tensor<4xf32>
+}
+{-#
+  dialect_resources: {
+    builtin: {
+      reshape_resource: "0x040000003f800000400000004040000040800000"
+    }
+  }
+#-}
+
+// -----
+
 // CHECK-LABEL: @slice_splat
 func.func @slice_splat() -> tensor<1x1x1xi32> {
   // CHECK: %[[SLICE:.+]] = "tosa.const"() <{values = dense<42> : tensor<1x1x1xi32>}

>From 4849dc0082adef2656b1e130c76b096e45618cae Mon Sep 17 00:00:00 2001
From: Thibaut Goetghebuer-Planchon <thibaut.goetghebuer-planchon at arm.com>
Date: Tue, 4 Aug 2026 10:50:53 +0100
Subject: [PATCH 2/2] [mlir][tosa] Deduplicate static shape check in ReshapeOp
 folding

---
 .../lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index d3ded54f5d909..d4ae8e920f489 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2117,23 +2117,18 @@ OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
   if (!inputTy.getElementType().isIntOrIndexOrFloat())
     return {};
 
+  // Constants must have static shape.
+  if (!outputTy.hasStaticShape())
+    return {};
+
   // Reshaping a resource-backed constant only requires updating its type.
   if (auto operand = llvm::dyn_cast_if_present<DenseResourceElementsAttr>(
-          adaptor.getInput1())) {
-    if (!outputTy.hasStaticShape()) {
-      return {};
-    }
-
+          adaptor.getInput1()))
     return DenseResourceElementsAttr::get(outputTy, operand.getRawHandle());
-  }
 
   // reshape(const(x)) -> const(reshape-attr(x))
   if (auto operand =
           llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1())) {
-    // Constants must have static shape.
-    if (!outputTy.hasStaticShape())
-      return {};
-
     // Okay to duplicate splat constants.
     if (operand.isSplat())
       return SplatElementsAttr::get(outputTy,