[mlir][tosa] Support `DenseResourceElementsAttr` in TOSA transpose folders #124532

GeorgeARM · 2025-01-27T11:41:16Z

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.

llvmbot · 2025-01-27T11:41:57Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Georgios Pinitas (GeorgeARM)

Changes

…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.

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

2 Files Affected:

(modified) mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp (+55)
(modified) mlir/test/Dialect/Tosa/constant-op-fold.mlir (+4-4)

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"
     }
   }
 #-}

sjarus · 2025-01-27T16:35:24Z

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

@@ -176,13 +177,38 @@ DenseElementsAttr transposeType(const RangeType &data, ShapedType inputType,
                                llvm::ArrayRef<ElementType>(outputValues));
 }

+// Function that tries to wrap the DenseResourceElementsAttr data access


Unclear wording.

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

Jerry-Ge · 2025-01-27T22:22:41Z

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

+        data && elementTy.isInteger(64))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<float>(attr);


are we missing other floating point types like FP16 per spec?

Needs handling through APFloat I presume. Need to look into this. Happy to do this as part of this patch if you want.

So had a look at this and is not as simple. I think this requires more effort to layer cleanly.

Is it worth adding this as a TODO below?

lhutton1

Small question otherwise LGTM!

lhutton1 · 2025-02-20T11:59:48Z

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

+        data && elementTy.isInteger(64))
+      return transposeType(*data, inputType, outputType, permValues);
+
+    if (auto data = tryGetDenseResourceValues<float>(attr);


Is it worth adding this as a TODO below?

Jerry-Ge · 2025-03-12T17:49:08Z

Hi @GeorgeARM , could you rebase again?

…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. Signed-off-by: Georgios Pinitas <[email protected]>

Jerry-Ge

LGTM

llvmbot added mlir mlir:tosa labels Jan 27, 2025

GeorgeARM requested review from sjarus, eric-k256 and TinaAMD January 27, 2025 11:41

sjarus reviewed Jan 27, 2025

View reviewed changes

GeorgeARM force-pushed the handle-dense branch from d7eb941 to a9663f2 Compare January 27, 2025 17:17

Jerry-Ge reviewed Jan 27, 2025

View reviewed changes

GeorgeARM mentioned this pull request Feb 11, 2025

[QUERY][MLIR] How to access dense resource #123961

Closed

GeorgeARM requested review from sjarus and lhutton1 February 13, 2025 13:57

lhutton1 approved these changes Feb 20, 2025

View reviewed changes

GeorgeARM force-pushed the handle-dense branch from 7daed89 to ea74bd6 Compare March 24, 2025 17:02

GeorgeARM force-pushed the handle-dense branch from ea74bd6 to 0f42343 Compare March 24, 2025 17:06

GeorgeARM requested a review from Jerry-Ge March 24, 2025 17:07

Jerry-Ge approved these changes Mar 24, 2025

View reviewed changes

GeorgeARM merged commit 3df9219 into llvm:main Mar 24, 2025
11 checks passed

GeorgeARM deleted the handle-dense branch May 2, 2025 13:00

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][tosa] Support `DenseResourceElementsAttr` in TOSA transpose folders #124532

[mlir][tosa] Support `DenseResourceElementsAttr` in TOSA transpose folders #124532

Uh oh!

GeorgeARM commented Jan 27, 2025 •

edited

Loading

Uh oh!

llvmbot commented Jan 27, 2025 •

edited

Loading

Uh oh!

sjarus Jan 27, 2025

Uh oh!

Uh oh!

Jerry-Ge Jan 27, 2025

Uh oh!

GeorgeARM Jan 29, 2025 •

edited

Loading

Uh oh!

GeorgeARM Feb 11, 2025

Uh oh!

lhutton1 Feb 20, 2025

Uh oh!

lhutton1 left a comment

Uh oh!

lhutton1 Feb 20, 2025

Uh oh!

Jerry-Ge commented Mar 12, 2025

Uh oh!

Jerry-Ge left a comment

Uh oh!

Uh oh!

Uh oh!

[mlir][tosa] Support DenseResourceElementsAttr in TOSA transpose folders #124532

[mlir][tosa] Support DenseResourceElementsAttr in TOSA transpose folders #124532

Uh oh!

Conversation

GeorgeARM commented Jan 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jan 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sjarus Jan 27, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Jerry-Ge Jan 27, 2025

Choose a reason for hiding this comment

Uh oh!

GeorgeARM Jan 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

GeorgeARM Feb 11, 2025

Choose a reason for hiding this comment

Uh oh!

lhutton1 Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

lhutton1 left a comment

Choose a reason for hiding this comment

Uh oh!

lhutton1 Feb 20, 2025

Choose a reason for hiding this comment

Uh oh!

Jerry-Ge commented Mar 12, 2025

Uh oh!

Jerry-Ge left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[mlir][tosa] Support `DenseResourceElementsAttr` in TOSA transpose folders #124532

[mlir][tosa] Support `DenseResourceElementsAttr` in TOSA transpose folders #124532

GeorgeARM commented Jan 27, 2025 •

edited

Loading

llvmbot commented Jan 27, 2025 •

edited

Loading

GeorgeARM Jan 29, 2025 •

edited

Loading