[Mlir-commits] [mlir] 7448e3e - [mlir][tosa] Support more float types in resource transpose folding (#213226)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 6 02:00:42 PDT 2026
Author: Thibaut Goetghebuer-Planchon
Date: 2026-08-06T10:00:36+01:00
New Revision: 7448e3e3af1a0a3a4cfc8846669085061177f54b
URL: https://github.com/llvm/llvm-project/commit/7448e3e3af1a0a3a4cfc8846669085061177f54b
DIFF: https://github.com/llvm/llvm-project/commit/7448e3e3af1a0a3a4cfc8846669085061177f54b.diff
LOG: [mlir][tosa] Support more float types in resource transpose folding (#213226)
Add support for f4, f8, f16, bf16, and f64 constants backed by
DenseResourceElementsAttr.
Added:
Modified:
mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp
mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
Removed:
################################################################################
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..296c72bb366a4 100644
--- a/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir
@@ -165,6 +165,126 @@ func.func @transpose_fold_dense_resource() -> tensor<2x2xf32> {
}
#-}
+// -----
+
+// CHECK-LABEL: @transpose_fold_dense_resource_f8e4m3fn
+func.func @transpose_fold_dense_resource_f8e4m3fn() -> tensor<2x2xf8E4M3FN> {
+ %0 = "tosa.const"() <{values = dense_resource<resource> : tensor<2x2xf8E4M3FN>}> : () -> tensor<2x2xf8E4M3FN>
+
+ // CHECK: %[[CST:.+]] = "tosa.const"() <{
+ // CHECK-SAME{LITERAL}: values = dense<[[1.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00]]> : tensor<2x2xf8E4M3FN>
+ %1 = tosa.transpose %0 { perms = array<i32: 1, 0> }: (tensor<2x2xf8E4M3FN>) -> tensor<2x2xf8E4M3FN>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xf8E4M3FN>
+}
+{-#
+ dialect_resources: {
+ builtin: {
+ resource: "0x0100000038404448"
+ }
+ }
+#-}
+
+// -----
+
+// 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"
+ }
+ }
+#-}
+
+// -----
+
+// CHECK-LABEL: @transpose_fold_dense_resource_f64
+func.func @transpose_fold_dense_resource_f64() -> tensor<2x2xf64> {
+ %0 = "tosa.const"() <{values = dense_resource<resource> : tensor<2x2xf64>}> : () -> tensor<2x2xf64>
+
+ // CHECK: %[[CST:.+]] = "tosa.const"() <{
+ // CHECK-SAME{LITERAL}: values = dense<[[1.000000e+00, 3.000000e+00], [2.000000e+00, 4.000000e+00]]> : tensor<2x2xf64>
+ %1 = tosa.transpose %0 { perms = array<i32: 1, 0> }: (tensor<2x2xf64>) -> tensor<2x2xf64>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xf64>
+}
+{-#
+ dialect_resources: {
+ builtin: {
+ resource: "0x08000000000000000000f03f000000000000004000000000000008400000000000001040"
+ }
+ }
+#-}
+
// -----
func.func @reduce_sum_constant() -> tensor<1x3xi32> {
More information about the Mlir-commits
mailing list