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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Cathal Corbett (catcor01)

<details>
<summary>Changes</summary>

Adds TOSA canonicalization support for folding `tosa.reshape(tosa.const)` when the constant uses `DenseResourceElementsAttr`.

The fold rematerializes a `tosa.const` with the reshape result type while preserving the existing dense resource handle, avoiding a runtime reshape for resource-backed constants.

Tests:
- Added TOSA canonicalization coverage for resource-backed constants.
- `git diff --check` passes locally.
- Full local `mlir-opt ... -canonicalize` could not be completed in this vendored checkout because the available built `mlir-opt` is older than current `llvm-project/main` and fails on unrelated newer TOSA ops such as `tosa.row_gather`.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+11) 
- (modified) mlir/test/Dialect/Tosa/canonicalize.mlir (+20) 


``````````diff
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"
+    }
+  }
+#-}

``````````

</details>


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