Add option max-tile-fold-size to limit the output size of tosa.tiles to constant fold

Signed-off-by: Rickert, Jonas <[email protected]>

Author: jorickert

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td

@@ -4,6 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// Modifications (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its
+// affiliates
+
 //===----------------------------------------------------------------------===//
 //
 // This file declares the passes for the TOSA Dialect in MLIR.
@@ -35,6 +38,9 @@ def TosaLayerwiseConstantFoldPass : Pass<"tosa-layerwise-constant-fold", "func::
     Option<"enableTileFolding", "enable-tile-folding", "bool",
             /*default=*/"false",
             "Enables folding of tosa.tile operation. May generate bigger constant values.">,
+    Option<"maxTileFoldSize", "max-tile-fold-size", "int",
+            /*default=*/"0",
+            "Maximum size in bytes of the constant to be generated when folding a tosa.tile operation. A value of 0 means no limit.">,
   ];
 }
 

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

@@ -1939,7 +1939,11 @@ DenseElementsAttr tile(DenseElementsAttr inputValues, ShapedType outputType) {
 }
 
 struct TosaFoldConstantTile : public TosaFoldConstantBase<tosa::TileOp> {
-  using TosaFoldConstantBase::TosaFoldConstantBase;
+
+  TosaFoldConstantTile(MLIRContext *ctxt, bool foldSplatOrSingleUseOnly,
+                       int maxSizeToFold)
+      : TosaFoldConstantBase<tosa::TileOp>(ctxt, foldSplatOrSingleUseOnly),
+        maxSizeToFold(maxSizeToFold) {}
 
   LogicalResult matchAndRewrite(tosa::TileOp op,
                                 PatternRewriter &rewriter) const override {
@@ -1958,10 +1962,22 @@ struct TosaFoldConstantTile : public TosaFoldConstantBase<tosa::TileOp> {
         foldSplatOrSingleUseOnly)
       return failure();
 
+    assert(maxSizeToFold >= 0 && "maxSizeToFold should be non-negative");
+    if (maxSizeToFold > 0) {
+      if (!outputType.hasStaticShape())
+        return failure();
+      const int64_t numOfElements = outputType.getNumElements();
+      if (numOfElements * (outputType.getElementTypeBitWidth() / 8) >
+          static_cast<int64_t>(maxSizeToFold))
+        return failure();
+    }
+
     rewriter.replaceOpWithNewOp<tosa::ConstOp>(op, outputType,
                                                tile(inputValues, outputType));
     return success();
   }
+
+  const int maxSizeToFold;
 };
 
 /// Getting the axes position of the element which is located
@@ -2275,7 +2291,8 @@ void mlir::tosa::populateTosaFoldConstantPatterns(
   patterns.add<TosaFoldConstantMatMul>(ctx, options.foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantConcat>(ctx, options.foldSplatOrSingleUseOnly);
   if (options.enableTileFolding)
-    patterns.add<TosaFoldConstantTile>(ctx, options.foldSplatOrSingleUseOnly);
+    patterns.add<TosaFoldConstantTile>(ctx, options.foldSplatOrSingleUseOnly,
+                                       options.maxTileFoldSize);
 }
 
 void mlir::tosa::populateTosaConstantReduction(MLIRContext *ctx,

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

@@ -4,6 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// Modifications (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its
+// affiliates
+//
 //===----------------------------------------------------------------------===//
 //
 // This file implements constant folding transformations on TOSA operations
@@ -57,6 +60,7 @@ struct TosaLayerwiseConstantFoldPass
     options.foldSplatOrSingleUseOnly = foldSplatOrSingleUseOnly;
     options.aggressiveReduceConstant = aggressiveReduceConstant;
     options.enableTileFolding = enableTileFolding;
+    options.maxTileFoldSize = maxTileFoldSize;
 
     mlir::tosa::populateTosaFoldConstantPatterns(ctx, patterns, options);
     mlir::tosa::populateTosaConstantReduction(ctx, patterns,

mlir/test/Dialect/Tosa/constant-tile.mlir

@@ -1,4 +1,6 @@
-// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold=enable-tile-folding=true %s | FileCheck %s
+// Modifications (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its
+// affiliates
+// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold="enable-tile-folding=true max-tile-fold-size=1024" %s | FileCheck %s
 // RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold %s | FileCheck --check-prefix=NO-FOLDING-CHECK %s
 
 // CHECK-LABEL: @tile_int_one_dim
@@ -12,6 +14,16 @@ func.func @tile_int_one_dim() -> (tensor<6xi32>) {
   // NO-FOLDING-CHECK: tosa.tile
 }
 
+// CHECK-LABEL: @tile_int_one_dim_large
+func.func @tile_int_one_dim_large() -> (tensor<1500xi32>) {
+  // CHECK: tosa.tile
+  %0 = "tosa.const"() {value = dense<[60, 2, 3]> : tensor<3xi32>} : () -> tensor<3xi32>
+  %cst = tosa.const_shape { value = dense<[500]> : tensor<1xindex> } : () -> !tosa.shape<1>
+  %1 = tosa.tile %0, %cst : (tensor<3xi32>, !tosa.shape<1>) -> tensor<1500xi32>
+  return %1 : tensor<1500xi32>
+  // NO-FOLDING-CHECK: tosa.tile
+}
+
 // CHECK-LABEL: @tile_bool
 func.func @tile_bool() -> (tensor<1x3x2x3xi1>) {
   // CHECK: "tosa.const"() <{value = dense