[Mlir-commits] [mlir] [mlir][tosa] Support `DenseResourceElementsAttr` in TOSA transpose folders (PR #124532)

[Georgios Pinitas]

https://github.com/GeorgeARM created https://github.com/llvm/llvm-project/pull/124532

…lder

Handle dense resource attributes in the transpose TOSA folder.
Currently their interface does not align with the rest of the `ElementsAttr` when it comes to data accessing hence the special handling.

>From d7eb941bb8d3c60b4cb82672241ddab52d1d4044 Mon Sep 17 00:00:00 2001
From: Georgios Pinitas <georgios.pinitas at arm.com>
Date: Sat, 25 Jan 2025 20:21:33 +0000
Subject: [PATCH] [mlir][tosa] Support `DenseResourceElementsAttr` in TOSA
 transpose folder

Handle dense resource attributes in the transpose TOSA folder.
Currently their interface does not align with the rest of the
`ElementsAttr` when it comes to data accessing hence the special handling.

Signed-off-by: Georgios Pinitas <georgios.pinitas at arm.com>
---
 .../Dialect/Tosa/Transforms/TosaFolders.cpp   | 55 +++++++++++++++++++
 mlir/test/Dialect/Tosa/constant-op-fold.mlir  |  8 +--
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
index 403ac48b915595..dfb23e3e586a93 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/APFloat.h"
@@ -176,6 +177,30 @@ DenseElementsAttr transposeType(const RangeType &data, ShapedType inputType,
                                 llvm::ArrayRef<ElementType>(outputValues));
 }
 
+// Function that tries to wrap the DenseResourceElementsAttr data access
+// handling as unfortunately at the moment don't share the same interface
+// with DenseElementsAttr
+template <typename T>
+std::optional<ArrayRef<T>> tryGetDenseResourceValues(ElementsAttr attr) {
+  if (auto denseResource = dyn_cast<DenseResourceElementsAttr>(attr)) {
+    // Check that the resource memory blob exists
+    AsmResourceBlob *blob = denseResource.getRawHandle().getBlob();
+    if (!blob)
+      return std::nullopt;
+
+    // Check that the data are in a valid form
+    bool isSplat = false;
+    if (!DenseElementsAttr::isValidRawBuffer(attr.getShapedType(),
+                                             blob->getData(), isSplat)) {
+      return std::nullopt;
+    }
+
+    return blob->template getDataAs<T>();
+  }
+
+  return std::nullopt;
+}
+
 // A type specialized transposition of an ElementsAttr.
 // This implementation tries to operate on the underlying data in its raw
 // representation when possible to avoid allocating a large number of Attribute
@@ -183,6 +208,7 @@ DenseElementsAttr transposeType(const RangeType &data, ShapedType inputType,
 DenseElementsAttr transpose(ElementsAttr attr, ShapedType inputType,
                             ShapedType outputType,
                             llvm::ArrayRef<int64_t> permValues) {
+  // Handle generic ElementsAttr
   if (auto data = attr.tryGetValues<bool>())
     return transposeType(*data, inputType, outputType, permValues);
 
@@ -204,6 +230,35 @@ DenseElementsAttr transpose(ElementsAttr attr, ShapedType inputType,
   if (auto data = attr.tryGetValues<APFloat>())
     return transposeType(*data, inputType, outputType, permValues);
 
+  // Handle DenseResourceElementsAttr
+  if (isa<DenseResourceElementsAttr>(attr)) {
+    auto elementTy = attr.getElementType();
+
+    if (auto data = tryGetDenseResourceValues<bool>(attr);
+        data && elementTy.isInteger(1))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<int8_t>(attr);
+        data && elementTy.isInteger(8))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<int16_t>(attr);
+        data && elementTy.isInteger(16))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<int32_t>(attr);
+        data && elementTy.isInteger(32))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<int64_t>(attr);
+        data && elementTy.isInteger(64))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<float>(attr);
+        data && elementTy.isF32())
+      return transposeType(*data, inputType, outputType, permValues);
+  }
+
   return nullptr;
 }
 
diff --git a/mlir/test/Dialect/Tosa/constant-op-fold.mlir b/mlir/test/Dialect/Tosa/constant-op-fold.mlir
index 8198903b78ac05..b1d0d2dc756b31 100644
--- a/mlir/test/Dialect/Tosa/constant-op-fold.mlir
+++ b/mlir/test/Dialect/Tosa/constant-op-fold.mlir
@@ -117,19 +117,19 @@ func.func @transpose_nofold_multi_users() -> (tensor<3x2xf32>, tensor<2x3xf32>)
   return %1, %input : tensor<3x2xf32>, tensor<2x3xf32>
 }
 
-// CHECK-LABEL: @transpose_nofold_dense_resource
-func.func @transpose_nofold_dense_resource() -> tensor<2x2xf32> {
+// CHECK-LABEL: @transpose_fold_dense_resource
+func.func @transpose_fold_dense_resource() -> tensor<2x2xf32> {
   %0 = "tosa.const"() <{value = dense_resource<resource> : tensor<2x2xf32>}> : () -> tensor<2x2xf32>
   %1 = "tosa.const"() <{value = dense<[1, 0]> : tensor<2xi32>}> : () -> tensor<2xi32>
 
-  // CHECK: tosa.transpose
+  // CHECK-NOT: tosa.transpose
   %2 = tosa.transpose %0, %1 : (tensor<2x2xf32>, tensor<2xi32>) -> tensor<2x2xf32>
   return %2 : tensor<2x2xf32>
 }
 {-#
   dialect_resources: {
     builtin: {
-      resource: "0x08000000010000000000000002000000000000000300000000000000"
+      resource: "0x040000003f800000400000004040000040800000"
     }
   }
 #-}