Repository: alibaba/BladeDISC (language: python)

Issue:
FileCheck does not enable TorchToMlir conversion pass



Discussion / hints:

Fix patch:
diff --git a/pytorch_blade/src/torch-mlir/BUILD b/pytorch_blade/src/torch-mlir/BUILD
index 22e64541e05..71e3df170aa 100644
--- a/pytorch_blade/src/torch-mlir/BUILD
+++ b/pytorch_blade/src/torch-mlir/BUILD
@@ -40,6 +40,7 @@ cc_library(
     name = "TorchMLIRTorchToMhlo",
     srcs = [
         "lib/Conversion/TorchToMhlo/TorchToMhlo.cpp",
+        "lib/Conversion/MhloPasses.cpp",
         "lib/Conversion/TorchToMhlo/MhloLegalizeUtils.cpp",
         "lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp",
         "lib/Dialect/TorchConversion/Transforms/ApplyValueSemantics.cpp",
@@ -47,7 +48,7 @@ cc_library(
     ] ,
     hdrs = [
         "include/torch-mlir/Conversion/MhloPasses.h",
-	"include/torch-mlir/Conversion/TorchToMhlo/ShapeUtils.h",
+	    "include/torch-mlir/Conversion/TorchToMhlo/ShapeUtils.h",
         "include/torch-mlir/Conversion/TorchToMhlo/TorchToMhlo.h",
         "include/torch-mlir/Conversion/TorchToMhlo/MhloLegalizeUtils.h"
     ],
diff --git a/pytorch_blade/src/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h b/pytorch_blade/src/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h
index 55d8ca20c24..798c0e3762a 100644
--- a/pytorch_blade/src/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h
+++ b/pytorch_blade/src/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h
@@ -34,6 +34,10 @@ namespace TorchConversion {
 #define GEN_PASS_CLASSES
 #include "torch-mlir/Conversion/MhloPasses.h.inc"
 } // namespace TorchConversion
+
+/// Registers all TorchToMhlo conversion passes.
+void registerTorchToMhloPasses();
+
 } // namespace torch
 } // end namespace mlir
 
diff --git a/pytorch_blade/src/torch-mlir/lib/Conversion/MhloPasses.cpp b/pytorch_blade/src/torch-mlir/lib/Conversion/MhloPasses.cpp
new file mode 100644
index 00000000000..6cac5a88a63
--- /dev/null
+++ b/pytorch_blade/src/torch-mlir/lib/Conversion/MhloPasses.cpp
@@ -0,0 +1,29 @@
+// Copyright 2022 The BladeDISC Authors. 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.
+
+#include "torch-mlir/Conversion/MhloPasses.h"
+#include "torch-mlir/Conversion/TorchToMhlo/TorchToMhlo.h"
+//#include "mlir/Dialect/TorchConversion/Transforms/Passes.h"
+#include "torch-mlir/Dialect/Torch/Transforms/Passes.h"
+
+namespace {
+#define GEN_PASS_REGISTRATION
+#include "torch-mlir/Conversion/MhloPasses.h.inc"
+} // end namespace
+
+void ::mlir::torch::registerTorchToMhloPasses() {
+  ::registerPasses();
+  ::mlir::PassPipelineRegistration<Torch::TorchLoweringPipelineOptions>(
+      "torch-backend-to-mhlo-backend-pipeline",
+      "Pipeline lowering torch backend contract to mhlo backend "
+      "contract.",
+      ::mlir::torch::TorchConversion::createTorchBackendToMhloBackendPipeline);
+}
\ No newline at end of file
diff --git a/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp b/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
index 89c70f619d2..725aecb42df 100644
--- a/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
+++ b/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
@@ -1630,7 +1630,6 @@ class ConvertTorchToMhlo
     // INSERT_ATENOP_PATTERN(AtenUnsqueezeOp);
     INSERT_ATENOP_PATTERN(AtenDropoutOp);
     INSERT_ATENOP_PATTERN(AtenViewOp);
-    // INSERT_ATENOP_PATTERN(AtenGeluOp);
     // INSERT_ATENOP_PATTERN(AtenGeluBackwardOp);
 #undef INSERT_ATENOP_PATTERN
 


Test patch:
diff --git a/pytorch_blade/tests/mhlo/elementwise.mlir b/pytorch_blade/tests/mhlo/elementwise.mlir
new file mode 100644
index 00000000000..e9da8b1b229
--- /dev/null
+++ b/pytorch_blade/tests/mhlo/elementwise.mlir
@@ -0,0 +1,40 @@
+// RUN: torch-mlir-opt <%s -convert-torch-to-mhlo -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL:   func @torch.aten.sigmoid(
+// CHECK-SAME:          %[[ARG:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+// CHECK:           %[[BUILTIN_ARG:.*]] = torch_c.to_builtin_tensor %[[ARG]] : !torch.vtensor<[?,?],f32> -> tensor<?x?xf32>
+// CHECK:           %[[CST:.*]] = "chlo.constant_like"(%[[BUILTIN_ARG]]) {value = 1.000000e+00 : f32} : (tensor<?x?xf32>) -> tensor<?x?xf32
+// CHECK:           %[[NEG_ARG:.*]] = mhlo.negate %[[BUILTIN_ARG]] : tensor<?x?xf32>
+// CHECK:           %[[EXPONTIAL_NEG:.*]] = mhlo.exponential %[[NEG_ARG]] : tensor<?x?xf32>
+// CHECK:           %[[ADD:.*]] = mhlo.add %[[EXPONTIAL_NEG]], %[[CST]] : tensor<?x?xf32>
+// CHECK:           %[[DIV:.*]] = mhlo.divide %[[CST]], %[[ADD]] : tensor<?x?xf32>
+// CHECK:           %[[RESULT:.*]] = torch_c.from_builtin_tensor %[[DIV]] : tensor<?x?xf32> -> !torch.vtensor<[?,?],f32>
+// CHECK:           return %[[RESULT]] : !torch.vtensor<[?,?],f32>
+func @torch.aten.sigmoid(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+    %0 = torch.aten.sigmoid %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
+    return %0 : !torch.vtensor<[?,?],f32>
+}
+
+
+// CHECK-LABEL: func @torch.aten.relu(
+// CHECK-SAME:          %[[ARG:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+// CHECK:             %[[CMP_GT:.*]] = "mhlo.compare"(%{{.*}} {compare_type = #mhlo<"comparison_type NOTYPE">, comparison_direction = #mhlo<"comparison_direction GT">} : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xi1>
+// CHECK:             %[[SELECT:.*]] = "mhlo.select"(%[[CMP_GT]], %{{.*}}) : (tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+func @torch.aten.relu(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+    %0 = torch.aten.relu %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor<[?,?],f32>
+    return %0 : !torch.vtensor<[?,?],f32>
+}
+
+// CHECK-LABEL: func @torch.aten.gelu(
+// CHECK-SAME:          %[[ARG:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+// CHECK:         %[[RSQRT:.*]] = mhlo.rsqrt %{{.*}} : tensor<?x?xf32>
+// CHECK:         %[[MUL0:.*]] = mhlo.multiply %{{.*}}, %{{.*}} : tensor<?x?xf32>
+// CHECK:         %[[ERF:.*]] = chlo.erf %[[MUL0]] : tensor<?x?xf32> -> tensor<?x?xf32>
+// CHECK:         %[[ADD:.*]] = mhlo.add %[[ERF]], %{{.*}} : tensor<?x?xf32>
+// CHECK:         %[[MUL1:.*]] = mhlo.multiply %[[ADD]], %{{.*}} : tensor<?x?xf32>
+// CHECK:         %[[MUL2:.*]] = mhlo.multiply %{{.*}}, %[[MUL1]] : tensor<?x?xf32>
+func @torch.aten.gelu(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+    %str = torch.constant.str "none"
+    %0 = torch.aten.gelu %arg0, %str : !torch.vtensor<[?,?],f32>, !torch.str -> !torch.vtensor<[?,?],f32>
+    return %0 : !torch.vtensor<[?,?],f32>
+}
\ No newline at end of file
diff --git a/pytorch_blade/tests/mhlo/ops.mlir b/pytorch_blade/tests/mhlo/ops.mlir
index 7a03c29057e..06c141133e5 100644
--- a/pytorch_blade/tests/mhlo/ops.mlir
+++ b/pytorch_blade/tests/mhlo/ops.mlir
@@ -1,4 +1,4 @@
-// RUN: torch-mlir-opt %s | torch-mlir-opt | FileCheck %s
+// RUN: torch-mlir-opt <%s -convert-torch-to-mhlo -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func @builtin_tensor_interop(
 func @builtin_tensor_interop(%arg0: tensor<*xf32>, %arg1: tensor<3x?xi8>, %arg2: !torch.vtensor<*,f32>, %arg3: !torch.vtensor<[3,?],si8>) {
diff --git a/pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp b/pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp
index 5dd395458f2..6aa5181e403 100644
--- a/pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp
+++ b/pytorch_blade/tests/mhlo/torch-mlir-opt/torch-mlir-opt.cpp
@@ -21,11 +21,13 @@
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllPasses.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
+#include "torch-mlir/Conversion/MhloPasses.h"
 #include "torch-mlir/InitAll.h"
 
 using namespace mlir;
 
 int main(int argc, char** argv) {
+  mlir::torch::registerTorchToMhloPasses();
   registerAllPasses();
   mlir::torch::registerAllPasses();