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

[Cathal Corbett]

https://github.com/catcor01 updated https://github.com/llvm/llvm-project/pull/212461

>From b72eda2b86327bb0be56b891c657daee1169c6cb Mon Sep 17 00:00:00 2001
From: Cathal Corbett <cathal.corbett at arm.com>
Date: Tue, 28 Jul 2026 12:36:18 +0200
Subject: [PATCH] [TOSA] Fold reshape of resource-backed constants

Teach tosa.reshape folding to rematerialize DenseResourceElementsAttr constants with the reshaped result type. Resource-backed constants can share the same blob handle, so this removes layout-only reshape ops without duplicating payload data.

Add canonicalization coverage for resource-backed tosa.const reshape folding.

Tests:

- mlir-opt mlir/test/Dialect/Tosa/canonicalize.mlir -canonicalize -split-input-file | FileCheck mlir/test/Dialect/Tosa/canonicalize.mlir

Change-Id: I63ec1ff2258b5e9676a7abd93a19d3bdfef9883b
---
 .../Dialect/Tosa/IR/TosaCanonicalizations.cpp | 11 ++++++++++
 mlir/test/Dialect/Tosa/canonicalize.mlir      | 20 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 8937875df7fc5..5ef5cf85be5b4 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -2140,6 +2140,17 @@ OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
         llvm::cast<ShapedType>(operand.getType()).clone(shapeVec));
   }
 
+  if (auto operand = llvm::dyn_cast_if_present<DenseResourceElementsAttr>(
+          adaptor.getInput1())) {
+    // Constants must have static shape.
+    if (!outputTy.hasStaticShape())
+      return {};
+
+    // Resource-backed constants can be rematerialized as a new view of the same
+    // underlying blob without duplicating the resource payload.
+    return DenseResourceElementsAttr::get(outputTy, operand.getRawHandle());
+  }
+
   return {};
 }
 
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 7c9dd260ca497..ecc4b94d49b61 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -759,6 +759,19 @@ func.func @reshape_canonicalize_const_splat() -> (tensor<10xi32>, tensor<1x10xi3
 
 // -----
 
+// CHECK-LABEL: @reshape_canonicalize_const_resource
+func.func @reshape_canonicalize_const_resource() -> (tensor<5xi32>, tensor<1x5xi32>) {
+  // CHECK-DAG: %[[VAR0:.+]] = "tosa.const"() <{values = dense_resource<reshape_resource> : tensor<5xi32>}
+  // CHECK-DAG: %[[VAR1:.+]] = "tosa.const"() <{values = dense_resource<reshape_resource> : tensor<1x5xi32>}
+  // CHECK: return %[[VAR0]], %[[VAR1]]
+  %0 = "tosa.const"() {values = dense_resource<reshape_resource> : tensor<5xi32>} : () -> tensor<5xi32>
+  %2 = "tosa.const_shape"() {values = dense<[1, 5]> : tensor<2xindex>} : () -> !tosa.shape<2>
+  %1 = tosa.reshape %0, %2 : (tensor<5xi32>, !tosa.shape<2>) -> tensor<1x5xi32>
+  return %0 , %1 : tensor<5xi32>, tensor<1x5xi32>
+}
+
+// -----
+
 // CHECK-LABEL: @reshape_canonicalize_const_sparse
 func.func @reshape_canonicalize_const_sparse() -> (tensor<3xi32>, tensor<1x3xi32>) {
   // CHECK: tosa.reshape
@@ -2266,3 +2279,10 @@ func.func @test_partially_foldable(%arg0: tensor<1x1x8x8xf32>, %arg1: tensor<1x2
   %2 = tosa.concat %0, %1 {axis = 1 : i32} : (tensor<1x2x8x8xf32>, tensor<1x2x8x8xf32>) -> tensor<1x4x8x8xf32>
   return %2 : tensor<1x4x8x8xf32>
 }
+{-#
+  dialect_resources: {
+    builtin: {
+      reshape_resource: "0x040000000000000001000000020000000300000004000000"
+    }
+  }
+#-}