[Mlir-commits] [mlir] [mlir][tosa] Support more float types in resource transpose folding (PR #213226)

[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 f4, f8, f16, bf16, and f64 constants backed by DenseResourceElementsAttr.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp (+40-4) 
- (modified) mlir/test/Dialect/Tosa/tosa-layerwise-constant-fold.mlir (+80) 


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

``````````

</details>


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