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

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,8 @@ std::unique_ptr<Pass> createTosaLayerwiseConstantFoldPass(
     const TosaLayerwiseConstantFoldPassOptions &options);
 std::unique_ptr<Pass> createTosaInferShapesPass();
 std::unique_ptr<Pass> createTosaMakeBroadcastablePass();
+std::unique_ptr<Pass> createSinkInputOpsThroughConcatPass(llvm::raw_ostream &);
+std::unique_ptr<Pass> createSinkInputOpsThroughConcatPass();
 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
 //
 //===----------------------------------------------------------------------===//
 //
@@ -142,4 +143,24 @@ 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".
+  }];
+}
+
 #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