Added folding

include/ttmlir/Dialect/TTIR/IR/TTIROps.td

@@ -899,6 +899,8 @@ def TTIR_ReshapeOp: TTIR_DPSOp<"reshape"> {
     }];
 
     let hasVerifier = 1;
+
+    let hasFolder = 1;
 }
 
 def TTIR_SliceOp: TTIR_DPSOp<"slice"> {

include/ttmlir/Dialect/TTIR/Transforms/Passes.td

@@ -112,11 +112,4 @@ def TTIRLoadSystemDesc: Pass<"ttir-load-system-desc", "::mlir::ModuleOp"> {
     ];
 }
 
-def TTIRReshapeFold: Pass<"ttir-reshape-fold", "::mlir::ModuleOp"> {
-  let summary = "Folds reshape ops that reshape the same shapes.";
-  let description = [{
-    This pass converts folds the ttir.reshape ops that are called to reshape a tensor to the same shape.
-  }];
-}
-
 #endif

lib/Dialect/TTIR/IR/TTIROps.cpp

@@ -358,6 +358,15 @@ ::mlir::LogicalResult mlir::tt::ttir::ReshapeOp::verify() {
   return success();
 }
 
+// ReshapeOp folder
+::mlir::OpFoldResult mlir::tt::ttir::ReshapeOp::fold(FoldAdaptor adaptor) {
+
+  if (getType() == getOperand(0).getType()) {
+    return getOperand(0);
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // SliceOp
 //===----------------------------------------------------------------------===//

lib/Dialect/TTIR/Transforms/CMakeLists.txt

@@ -3,7 +3,6 @@ add_mlir_dialect_library(MLIRTTIRTransforms
         Constant.cpp
         Generic.cpp
         Layout.cpp
-        Reshape.cpp
         Transforms.cpp
         Utility.cpp
 

lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp

@@ -72,11 +72,6 @@ void createTTNNPipelineDeallocPass(
   pm.addPass(createTTNNDeallocate());
 }
 
-void createTTIRReshapeFoldPass(
-    OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
-  pm.addPass(mlir::tt::ttir::createTTIRReshapeFold());
-}
-
 void createTTNNPipelineTTIRPassesFromString(OpPassManager &pm,
                                             std::string options) {
   auto optionsStruct =
@@ -112,16 +107,8 @@ void createTTNNPipelineDeallocPassFromString(OpPassManager &pm,
   createTTNNPipelineDeallocPass(pm, *optionsStruct);
 }
 
-void createTTIRReshapeFoldPassFromString(OpPassManager &pm,
-                                         std::string options) {
-  auto optionsStruct =
-      TTIRToTTNNBackendPipelineOptions::createFromString(options);
-  createTTNNPipelineDeallocPass(pm, *optionsStruct);
-}
-
 void createTTIRToTTNNBackendPipeline(
     OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options) {
-  createTTIRReshapeFoldPass(pm, options);
   createTTNNPipelineTTIRPasses(pm, options);
   createTTNNPipelineLoweringPasses(pm, options);
   createTTNNPipelineAnalysisPasses(pm, options);

test/ttmlir/Dialect/TTIR/reshape/reshape_test.mlir → test/ttmlir/Dialect/TTNN/reshape/reshape_folding_test.mlir

@@ -1,11 +1,11 @@
-// RUN: ttmlir-opt --ttir-reshape-fold %s| FileCheck %s
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s| FileCheck %s
 #any_device_tile = #tt.operand_constraint<dram|l1|tile|any_device_tile>
 // Tests if we fold when translating from "ttir.reshape" which is called on the two same shapes.
 module @reshape_test {
   func.func @main(%arg0: tensor<1xi32>) -> (tensor<1xi32> {jax.result_info = ""}) {
     %0 = tensor.empty() : tensor<1xi32>
     %1 = "ttir.reshape"(%arg0, %0) <{operand_constraints = [#any_device_tile, #any_device_tile], shape = [1 : i32]}> : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
-    // CHECK: return %arg0 : tensor<1xi32>
+    // CHECK: return %arg0 : tensor<1xi32, #{{.*}}>
     return %1 : tensor<1xi32>
   }
 }