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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tosa

Author: Thibaut Goetghebuer-Planchon (Tessil)

<details>
<summary>Changes</summary>

Add support for `DenseResourceElementsAttr` in TOSA reshape folding.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp (+10) 
- (modified) mlir/test/Dialect/Tosa/constant_folding.mlir (+19) 


``````````diff
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>}

``````````

</details>


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