[Mlir-commits] [mlir] [mlir][tosa] Support more float types in resource transpose folding (PR #213226)
Thibaut Goetghebuer-Planchon
llvmlistbot at llvm.org
Fri Jul 31 01:34:15 PDT 2026
https://github.com/Tessil created https://github.com/llvm/llvm-project/pull/213226
Add support for f4, f8, f16, bf16, and f64 constants backed by DenseResourceElementsAttr.
>From 5b5bfb4310bd7a22b4620514d89a3be78c20f0cb Mon Sep 17 00:00:00 2001
From: Thibaut Goetghebuer-Planchon <thibaut.goetghebuer-planchon at arm.com>
Date: Fri, 31 Jul 2026 09:32:31 +0100
Subject: [PATCH] [mlir][tosa] Support more float types in resource transpose
folding
Add support for f4, f8, f16, bf16, and f64 constants backed by
DenseResourceElementsAttr.
---
.../Dialect/Tosa/Transforms/TosaFolders.cpp | 44 +++++++++-
.../Tosa/tosa-layerwise-constant-fold.mlir | 80 +++++++++++++++++++
2 files changed, 120 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
index d6961628afc9f..604965f7103b1 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
@@ -131,15 +131,15 @@ bool constantUnaryOpShouldBeFolded(TosaOp unaryOp, DenseElementsAttr values) {
}
template <typename RangeType>
-DenseElementsAttr transposeType(const RangeType &data, ShapedType inputType,
- ShapedType outputType,
- llvm::ArrayRef<int64_t> permValues) {
+auto transposeValues(const RangeType &data, ShapedType inputType,
+ ShapedType outputType,
+ llvm::ArrayRef<int64_t> permValues) {
using ElementType = std::decay_t<decltype(*std::begin(data))>;
assert(inputType.getElementType() == outputType.getElementType());
if (inputType.getNumElements() == 0)
- return DenseElementsAttr::get(outputType, llvm::ArrayRef<ElementType>{});
+ return SmallVector<ElementType>{};
auto inputShape = inputType.getShape();
@@ -171,10 +171,33 @@ DenseElementsAttr transposeType(const RangeType &data, ShapedType inputType,
outputValues[dstLinearIndex] = it.value();
}
+ return outputValues;
+}
+
+template <typename RangeType>
+DenseElementsAttr transposeType(const RangeType &data, ShapedType inputType,
+ ShapedType outputType,
+ llvm::ArrayRef<int64_t> permValues) {
+ using ElementType = std::decay_t<decltype(*std::begin(data))>;
+ SmallVector<ElementType> outputValues =
+ transposeValues(data, inputType, outputType, permValues);
return DenseElementsAttr::get(outputType,
llvm::ArrayRef<ElementType>(outputValues));
}
+template <typename RangeType>
+DenseElementsAttr transposeRawType(const RangeType &data, ShapedType inputType,
+ ShapedType outputType,
+ llvm::ArrayRef<int64_t> permValues) {
+ using StorageType = std::decay_t<decltype(*std::begin(data))>;
+ SmallVector<StorageType> outputValues =
+ transposeValues(data, inputType, outputType, permValues);
+ llvm::ArrayRef<char> rawData(
+ reinterpret_cast<const char *>(outputValues.data()),
+ outputValues.size() * sizeof(StorageType));
+ return DenseElementsAttr::getFromRawBuffer(outputType, rawData);
+}
+
// 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
@@ -231,6 +254,19 @@ DenseElementsAttr transpose(ElementsAttr attr, ShapedType inputType,
if (auto data = tryGetDenseResourceValues<float>(attr);
data && elementTy.isF32())
return transposeType(*data, inputType, outputType, permValues);
+
+ if (auto data = tryGetDenseResourceValues<uint8_t>(attr);
+ data && isa<Float4E2M1FNType, Float8E4M3FNType, Float8E5M2Type,
+ Float8E8M0FNUType>(elementTy))
+ return transposeRawType(*data, inputType, outputType, permValues);
+
+ if (auto data = tryGetDenseResourceValues<uint16_t>(attr);
+ data && isa<Float16Type, BFloat16Type>(elementTy))
+ return transposeRawType(*data, inputType, outputType, permValues);
+
+ if (auto data = tryGetDenseResourceValues<double>(attr);
+ data && elementTy.isF64())
+ return transposeType(*data, inputType, outputType, permValues);
}
return nullptr;
diff --git a/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir b/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
index 711dfe4d2405e..9290f088c87c9 100644
--- a/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
@@ -165,6 +165,86 @@ func.func @transpose_fold_dense_resource() -> tensor<2x2xf32> {
}
#-}
+// -----
+
+// CHECK-LABEL: @transpose_fold_dense_resource_f8e5m2
+func.func @transpose_fold_dense_resource_f8e5m2() -> tensor<2x2xf8E5M2> {
+ %0 = "tosa.const"() <{values = dense_resource<resource> : tensor<2x2xf8E5M2>}> : () -> tensor<2x2xf8E5M2>
+
+ // CHECK: %[[CST:.+]] = "tosa.const"() <{
+ // CHECK-SAME{LITERAL}: values = dense<[[1.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00]]> : tensor<2x2xf8E5M2>
+ %1 = tosa.transpose %0 { perms = array<i32: 1, 0> }: (tensor<2x2xf8E5M2>) -> tensor<2x2xf8E5M2>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xf8E5M2>
+}
+{-#
+ dialect_resources: {
+ builtin: {
+ resource: "0x010000003c404244"
+ }
+ }
+#-}
+
+// -----
+
+// CHECK-LABEL: @transpose_fold_dense_resource_f4e2m1fn
+func.func @transpose_fold_dense_resource_f4e2m1fn() -> tensor<2x2xf4E2M1FN> {
+ %0 = "tosa.const"() <{values = dense_resource<resource> : tensor<2x2xf4E2M1FN>}> : () -> tensor<2x2xf4E2M1FN>
+
+ // CHECK: %[[CST:.+]] = "tosa.const"() <{
+ // CHECK-SAME{LITERAL}: values = dense<[[1.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00]]> : tensor<2x2xf4E2M1FN>
+ %1 = tosa.transpose %0 { perms = array<i32: 1, 0> }: (tensor<2x2xf4E2M1FN>) -> tensor<2x2xf4E2M1FN>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xf4E2M1FN>
+}
+{-#
+ dialect_resources: {
+ builtin: {
+ resource: "0x0100000002040506"
+ }
+ }
+#-}
+
+// -----
+
+// CHECK-LABEL: @transpose_fold_dense_resource_f16
+func.func @transpose_fold_dense_resource_f16() -> tensor<2x2xf16> {
+ %0 = "tosa.const"() <{values = dense_resource<resource> : tensor<2x2xf16>}> : () -> tensor<2x2xf16>
+
+ // CHECK: %[[CST:.+]] = "tosa.const"() <{
+ // CHECK-SAME{LITERAL}: values = dense<[[1.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00]]> : tensor<2x2xf16>
+ %1 = tosa.transpose %0 { perms = array<i32: 1, 0> }: (tensor<2x2xf16>) -> tensor<2x2xf16>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xf16>
+}
+{-#
+ dialect_resources: {
+ builtin: {
+ resource: "0x02000000003c004000420044"
+ }
+ }
+#-}
+
+// -----
+
+// CHECK-LABEL: @transpose_fold_dense_resource_bf16
+func.func @transpose_fold_dense_resource_bf16() -> tensor<2x2xbf16> {
+ %0 = "tosa.const"() <{values = dense_resource<resource> : tensor<2x2xbf16>}> : () -> tensor<2x2xbf16>
+
+ // CHECK: %[[CST:.+]] = "tosa.const"() <{
+ // CHECK-SAME{LITERAL}: values = dense<[[1.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00]]> : tensor<2x2xbf16>
+ %1 = tosa.transpose %0 { perms = array<i32: 1, 0> }: (tensor<2x2xbf16>) -> tensor<2x2xbf16>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xbf16>
+}
+{-#
+ dialect_resources: {
+ builtin: {
+ resource: "0x02000000803f004040408040"
+ }
+ }
+#-}
+
// -----
func.func @reduce_sum_constant() -> tensor<1x3xi32> {
More information about the Mlir-commits
mailing list