[torch-frontend] add decompose-on-torch pass

frontends/torch-frontend/torch-frontend/include/torch-frontend/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(Conversion)
 add_subdirectory(Transforms)
+add_subdirectory(Dialect/Torch/Transforms)

frontends/torch-frontend/torch-frontend/include/torch-frontend/Dialect/Torch/Transforms/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name TorchFrontendTorchTransforms)
+add_public_tablegen_target(TorchFrontendTorchTransformsPassIncGen)

frontends/torch-frontend/torch-frontend/include/torch-frontend/Dialect/Torch/Transforms/DecomposeOnTorch.h

@@ -0,0 +1,33 @@
+//===- DecomposeOnTorch.h -------------------------------------*--- C++ -*-===//
+//
+// Copyright 2022 ByteDance Ltd. and/or its affiliates. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCH_FRONTEND_TORCH_TRANSFORMS_DECOMPOSEONTORCH_H
+#define TORCH_FRONTEND_TORCH_TRANSFORMS_DECOMPOSEONTORCH_H
+
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace mlir {
+namespace func {
+class FuncOp;
+} // namespace func
+
+std::unique_ptr<OperationPass<func::FuncOp>> createDecomposeOnTorch();
+
+} // namespace mlir
+
+#endif // TORCH_FRONTEND_TORCH_TRANSFORMS_DECOMPOSEONTORCH_H

frontends/torch-frontend/torch-frontend/include/torch-frontend/Transforms/FuseOpOnTorch.h → frontends/torch-frontend/torch-frontend/include/torch-frontend/Dialect/Torch/Transforms/FuseOpOnTorch.h

@@ -15,8 +15,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef TORCH_FRONTEND_CONVERSION_FUSEOPONTORCH_H
-#define TORCH_FRONTEND_CONVERSION_FUSEOPONTORCH_H
+#ifndef TORCH_FRONTEND_TORCH_TRANSFORMS_FUSEOPONTORCH_H
+#define TORCH_FRONTEND_TORCH_TRANSFORMS_FUSEOPONTORCH_H
 
 #include "mlir/Pass/Pass.h"
 #include <memory>
@@ -30,4 +30,4 @@ std::unique_ptr<OperationPass<func::FuncOp>> createFuseOpOnTorch();
 
 } // namespace mlir
 
-#endif // TORCH_FRONTEND_CONVERSION_FUSEOPONTORCH_H
+#endif // TORCH_FRONTEND_TORCH_TRANSFORMS_FUSEOPONTORCH_H

frontends/torch-frontend/torch-frontend/include/torch-frontend/Dialect/Torch/Transforms/Passes.h

@@ -0,0 +1,33 @@
+//===- Passes.h ----------------------------------------------*--- C++ -*-===//
+//
+// Copyright 2022 ByteDance Ltd. and/or its affiliates. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TORCH_FRONTEND_TORCH_TRANSFORMS_PASSES
+#define TORCH_FRONTEND_TORCH_TRANSFORMS_PASSES
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "torch-frontend/Dialect/Torch/Transforms/DecomposeOnTorch.h"
+#include "torch-frontend/Dialect/Torch/Transforms/FuseOpOnTorch.h"
+
+namespace mlir {
+
+// Generate the code for registering transforms passes.
+#define GEN_PASS_REGISTRATION
+#include "torch-frontend/Dialect/Torch/Transforms/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // TORCH_FRONTEND_TORCH_TRANSFORMS_PASSES

frontends/torch-frontend/torch-frontend/include/torch-frontend/Dialect/Torch/Transforms/Passes.td

@@ -0,0 +1,24 @@
+#ifndef TORCH_FRONTEND_TORCH_TRANSFORMS_PASSES_TD
+#define TORCH_FRONTEND_TORCH_TRANSFORMS_PASSES_TD
+
+include "mlir/Pass/PassBase.td"
+
+//===----------------------------------------------------------------------===//
+// FuseOpOnTorch
+//===----------------------------------------------------------------------===//
+
+def FuseOpOnTorch : Pass<"fuse-op-on-torch", "func::FuncOp"> {
+  let summary = "fuse op on torch dialect";
+  let constructor = "mlir::createFuseOpOnTorch()";
+}
+
+//===----------------------------------------------------------------------===//
+// DecomposeOnTorch
+//===----------------------------------------------------------------------===//
+
+def DecomposeOnTorch : Pass<"decompose-on-torch", "func::FuncOp"> {
+    let summary = "decompose op on torch dialect";
+    let constructor = "mlir::createDecomposeOnTorch()";
+}
+
+#endif // TORCH_FRONTEND_TORCH_TRANSFORMS_PASSES_TD

frontends/torch-frontend/torch-frontend/include/torch-frontend/Transforms/Passes.h

@@ -21,7 +21,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "torch-frontend/Transforms/CanonicalizeExt.h"
 #include "torch-frontend/Transforms/EliminateUselessOp.h"
-#include "torch-frontend/Transforms/FuseOpOnTorch.h"
 #include "torch-frontend/Transforms/RewriteCustomOp.h"
 #include "torch-frontend/Transforms/RewriteEntryFuncName.h"
 #include "torch-frontend/Transforms/UnpackPublicFunctionReturn.h"

frontends/torch-frontend/torch-frontend/include/torch-frontend/Transforms/Passes.td

@@ -49,15 +49,6 @@ def EliminateUselessOp : Pass<"eliminate-useless-op", "func::FuncOp"> {
   let constructor = "mlir::createEliminateUselessOpPass()";
 }
 
-//===----------------------------------------------------------------------===//
-// FuseOpOnTorch
-//===----------------------------------------------------------------------===//
-
-def FuseOpOnTorch : Pass<"fuse-op-on-torch", "func::FuncOp"> {
-  let summary = "fuse op on torch dialect";
-  let constructor = "mlir::createFuseOpOnTorch()";
-}
-
 //===----------------------------------------------------------------------===//
 // RewriteEntryFuncName
 //===----------------------------------------------------------------------===//

frontends/torch-frontend/torch-frontend/lib/CAPI/Passes.cpp

@@ -16,6 +16,7 @@
 //===----------------------------------------------------------------------===//
 #include "torch-frontend-c/Passes.h"
 #include "torch-frontend/Conversion/Passes.h"
+#include "torch-frontend/Dialect/Torch/Transforms/Passes.h"
 #include "torch-frontend/Pipelines/Pipelines.h"
 #include "torch-frontend/Transforms/Passes.h"
 
@@ -30,5 +31,6 @@ void torchFrontendRegisterConversionPasses() {
 }
 
 void torchFrontendRegisterTransformsPasses() {
+  mlir::registerTorchFrontendTorchTransformsPasses();
   mlir::registerTorchFrontendTransformsPasses();
 }

frontends/torch-frontend/torch-frontend/lib/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_subdirectory(CAPI)
 add_subdirectory(Conversion)
+add_subdirectory(Dialect/Torch/Transforms)
 add_subdirectory(Pipelines)
 add_subdirectory(Transforms)
 add_subdirectory(Utils)

frontends/torch-frontend/torch-frontend/lib/Dialect/Torch/Transforms/CMakeLists.txt

@@ -0,0 +1,21 @@
+set(LLVM_TARGET_DEFINITIONS FuseOpOnTorchPattern.td)
+mlir_tablegen(FuseOpOnTorchPattern.inc -gen-rewriters)
+add_public_tablegen_target(FuseOpOnTorchPatternIncGen)
+
+add_mlir_library(TorchFrontendTorchTransforms
+  DecomposeOnTorch.cpp
+  FuseOpOnTorch.cpp
+
+  DEPENDS
+  TorchFrontendTorchTransformsPassIncGen
+  FuseOpOnTorchPatternIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRPass
+  MLIRDialect
+  TorchMLIRTorchDialect
+  TorchMLIRTorchUtils
+  TorchFrontendUtils
+)
+target_include_directories(TorchFrontendTorchTransforms PUBLIC ${CMAKE_CURRENT_BINARY_DIR})