feat(tosa): sink tosa ops through tosa.concat (#671)

Author: gerion-amd

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

@@ -3,6 +3,8 @@
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Modifications (c) Copyright 2025 Advanced Micro Devices, Inc. or its
+// affiliates
 //
 //===----------------------------------------------------------------------===//
 //
@@ -44,6 +46,9 @@ std::unique_ptr<Pass> createTosaLayerwiseConstantFoldPass(
     const TosaLayerwiseConstantFoldPassOptions &options);
 std::unique_ptr<Pass> createTosaInferShapesPass();
 std::unique_ptr<Pass> createTosaMakeBroadcastablePass();
+std::unique_ptr<Pass>
+createSinkInputOpsThroughConcatPass(SinkInputOpsThroughConcatOptions &,
+                                    llvm::raw_ostream &);
 std::unique_ptr<Pass> createTosaTestQuantUtilAPIPass();
 std::unique_ptr<Pass> createTosaOptionalDecompositions();
 

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

@@ -3,6 +3,7 @@
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Modifications (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
 //
 // Modifications (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its
 // affiliates
@@ -148,4 +149,41 @@ def TosaReduceTransposes : Pass<"tosa-reduce-transposes", "func::FuncOp"> {
   }];
 }
 
+def SinkInputOpsThroughConcat : Pass<"sink-input-ops-through-concat", "mlir::ModuleOp"> {
+  let summary = "Sinks same operation through a Concat operation";
+  let description = [{
+    Pass that sinks the same operation through a concatenation, simplifying
+    later optimizations.
+
+    To explain with a picture:
+    ```
+    Replacing            with
+    - Op -\              -\
+    - Op ---> Concat ==> ---> Concat -> Op
+    - Op -/              -/
+    ```
+
+    The pass works greedy (i.e., it sinks operator chains) and does not do any
+    cost-benefit assessment. It it restricted to an explicit list of tosa
+    operations as input for the concatentation that are known to be "sinkable".
+  }];
+  let options = [
+      Option<"enableElementwises", "enable-elementwises", "bool",
+             /*default=*/"true",
+             "Sink down elementwise operations.">,
+      Option<"enableReductions", "enable-reductions", "bool",
+             /*default=*/"true",
+             "Sink down reduce operations.">,
+      Option<"enableReshape", "enable-reshape", "bool",
+             /*default=*/"true",
+             "Sink down the reshape operation.">,
+      Option<"enableMatmul", "enable-matmul", "bool",
+             /*default=*/"true",
+             "Sink down the matrix multiplication.">,
+      Option<"matchUntransformedOperations", "match-untransformed-operations", "bool",
+             /*default=*/"false",
+             "Print operations that have potential to be sunken but are not actually sunken down.">,
+   ];
+}
+
 #endif // MLIR_DIALECT_TOSA_TRANSFORMS_PASSES

mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt

@@ -1,3 +1,5 @@
+# Modifications (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
+
 add_mlir_dialect_library(MLIRTosaTransforms
   TosaDecomposeTransposeConv.cpp
   TosaDecomposeDepthwise.cpp
@@ -9,6 +11,7 @@ add_mlir_dialect_library(MLIRTosaTransforms
   TosaReduceTransposes.cpp
   TosaTypeConverters.cpp
   TosaValidation.cpp
+  SinkInputOpsThroughConcat.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa/Transforms