Repository: alibaba/BladeDISC (language: python)

Issue:
Need a white list for torch-mlir converter



Discussion / hints:

Fix patch:
diff --git a/pytorch_blade/ci_build/build_pytorch_blade.sh b/pytorch_blade/ci_build/build_pytorch_blade.sh
index 694c3615882..90988b5713c 100755
--- a/pytorch_blade/ci_build/build_pytorch_blade.sh
+++ b/pytorch_blade/ci_build/build_pytorch_blade.sh
@@ -52,10 +52,7 @@ function ci_build() {
     # The following are UNIT TESTS
     export TORCH_BLADE_DEBUG_LOG=ON
     pytest tests tests -v -m "not ltc" 2>&1 | tee -a py_test.out
-    TORCH_DISC_USE_TORCH_MLIR=true python3 tests/disc/ops/test_activation.py
-    TORCH_DISC_USE_TORCH_MLIR=true python3 tests/disc/ops/test_unary_ops.py
-    TORCH_DISC_USE_TORCH_MLIR=true python3 tests/disc/ops/test_binary_ops.py
-    TORCH_DISC_USE_TORCH_MLIR=true python3 tests/disc/ops/test_broadcast.py
+    TORCH_DISC_USE_TORCH_MLIR=true pytest tests/disc/ops/ -v -k "TestDiscActivation or TestDiscUnaryOps or TestDiscBinaryOps or TestDiscBroadcast"
     # DEBUG=1 will trigger debug mode compilation
     DEBUG=1 python3 setup.py cpp_test 2>&1 | tee -a cpp_test.out;
     python3 setup.py bdist_wheel;
diff --git a/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.cpp b/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.cpp
index ae1c46272d5..944edb11b18 100644
--- a/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.cpp
+++ b/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.cpp
@@ -23,6 +23,7 @@
 #include "compiler/jit/tool_funcs.h"
 #include "compiler/mlir/converters/mhlo_converter_register.h"
 #include "compiler/mlir/converters/mlir_type_utils.h"
+#include "compiler/mlir/converters/torch_mlir_op_filter.h"
 
 #include "function_importer.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -200,6 +201,9 @@ void RegisterDialects(mlir::DialectRegistry& registry) {
 }
 
 bool IsMlirMhloSupported(const torch::jit::Node& node) {
+  if (IsTorchMlirAvailable()) {
+    return IsTorchMlirSupported(node);
+  }
   c10::optional<OpConverter> converter = GetMlirMhloConverter(node);
   try {
     if (converter) {
@@ -404,7 +408,7 @@ ConvertTorchToMhlo(std::shared_ptr<torch::jit::Graph> graph) {
 std::tuple<std::string, std::string, std::string, std::string>
 ConvertTorchScriptToMhlo(std::shared_ptr<torch::jit::Graph> graph) {
   try {
-    if (env::ReadBoolFromEnvVar("TORCH_DISC_USE_TORCH_MLIR", false)) {
+    if (IsTorchMlirAvailable()) {
       return ConvertTorchToMhlo(graph);
     } else {
       mlir::DialectRegistry registry;
@@ -422,5 +426,9 @@ ConvertTorchScriptToMhlo(std::shared_ptr<torch::jit::Graph> graph) {
   }
 }
 
+bool IsTorchMlirAvailable() {
+  return env::ReadBoolFromEnvVar("TORCH_DISC_USE_TORCH_MLIR", false);
+}
+
 } // namespace blade
 } // namespace torch
diff --git a/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.h b/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.h
index 2af69d24663..c774eb33fb3 100644
--- a/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.h
+++ b/pytorch_blade/src/compiler/mlir/converters/mhlo_conversion.h
@@ -22,6 +22,8 @@ class Node;
 namespace torch {
 namespace blade {
 
+bool IsTorchMlirAvailable();
+
 bool IsMlirMhloSupported(const torch::jit::Node&);
 
 // Return a pair of the MLIR module strings, with the first one in
diff --git a/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.cpp b/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.cpp
new file mode 100644
index 00000000000..899bbb01145
--- /dev/null
+++ b/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.cpp
@@ -0,0 +1,52 @@
+// 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 "compiler/mlir/converters/torch_mlir_op_filter.h"
+#include <torch/script.h>
+#include <string>
+#include <unordered_set>
+#include "compiler/mlir/converters/mhlo_conversion_context.h"
+
+namespace torch {
+namespace blade {
+std::unordered_set<std::string> GetTorchMlirWhiteList();
+
+bool IsTorchMlirSupported(const torch::jit::Node& node) {
+  auto schema = node.maybeSchema();
+  if (schema) {
+    // auto name = schema->operator_name().name;
+    return GetTorchMlirWhiteList().find(schema->operator_name().name) !=
+        GetTorchMlirWhiteList().end();
+  } else if (node.kind().is_prim()) {
+    auto name = c10::OperatorName(node.kind().toQualString(), "").name;
+    return GetTorchMlirWhiteList().find(name) != GetTorchMlirWhiteList().end();
+  }
+  return false;
+}
+
+std::unordered_set<std::string> GetTorchMlirWhiteList() {
+  return std::unordered_set<std::string>{
+      "aten::relu",     "aten::relu6",       "aten::leaky_relu",
+      "aten::sigmoid",  "aten::glu",         "aten::contiguous",
+      "aten::erf",      "aten::exp",         "aten::neg",
+      "aten::rsqrt",    "aten::neg",         "aten::to.dtype",
+      "aten::type_as",  "aten::bitwise_not", "aten::to",
+      "aten::gelu",     "aten::silu",        "aten::size",
+      "aten::sub",      "aten::add",         "aten::div",
+      "aten::mul",      "aten::eq",          "aten::lt",
+      "aten::ne",       "aten::__and__",     "aten::floor_divide",
+      "prim::Constant", "aten::expand_as",   "aten::repeat",
+      "aten::expand",
+  };
+}
+
+} //  namespace blade
+} //  namespace torch
\ No newline at end of file
diff --git a/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.h b/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.h
new file mode 100644
index 00000000000..47cd4a99cd6
--- /dev/null
+++ b/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.h
@@ -0,0 +1,23 @@
+// 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.
+
+#pragma once
+#include "compiler/mlir/converters/mhlo_conversion_context.h"
+
+namespace torch {
+namespace blade {
+
+// TODO(Yancey1989): should dry run torchToMhlo pass to
+// evaluate TorchScript symbol
+bool IsTorchMlirSupported(const torch::jit::Node& node);
+
+} //  namespace blade
+} //  namespace torch
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 a1b72e798b3..6911a970690 100644
--- a/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
+++ b/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
@@ -659,7 +659,7 @@ LogicalResult ConvertAtenOp<AtenReluOp>::matchAndRewrite(
     ConversionPatternRewriter& rewriter) const {
   Location loc = op.getLoc();
   Value input = adaptor.self();
-  auto inputTy = input.getType().template cast<RankedTensorType>();
+  auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
   if (!inputTy) {
     return op.emitError("Only RankedTensorType is supported in Aten ReLU op.");
   }
@@ -671,6 +671,28 @@ LogicalResult ConvertAtenOp<AtenReluOp>::matchAndRewrite(
   return success();
 }
 
+// Convert a Aten::Relu6 to HLO
+// Relu6(x) = min(AtenRelu(x), 6)
+template <>
+LogicalResult ConvertAtenOp<AtenRelu6Op>::matchAndRewrite(
+    AtenRelu6Op op,
+    OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  Location loc = op.getLoc();
+  Value input = adaptor.self();
+  auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
+  if (!inputTy) {
+    return op.emitError("Only RankedTensorType is supported in Aten ReLU op.");
+  }
+  Value relu = rewriter.create<AtenReluOp>(loc, inputTy, input);
+  Value six = chlo::getConstantLike(rewriter, loc, 6.0, input);
+  Value compareLtSix = rewriter.create<mhlo::CompareOp>(
+      loc, input, six, mhlo::ComparisonDirection::LT);
+  rewriter.replaceOpWithNewOp<mhlo::SelectOp>(
+      op, inputTy, compareLtSix, relu, six);
+  return success();
+}
+
 // Convert a Aten::LeakyRelu to HLO
 // LeakyRelu(x) = max(0, x) + negative_slop * min(0, x)
 template <>
@@ -681,7 +703,7 @@ LogicalResult ConvertAtenOp<AtenLeakyReluOp>::matchAndRewrite(
   Location loc = op.getLoc();
   Value input = adaptor.self();
   Value negativeSlope = op.negative_slope();
-  auto inputTy = input.getType().template cast<RankedTensorType>();
+  auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
   if (!inputTy) {
     return op.emitError(
         "Only RankedTensorType is supported in Aten LeakyReLU op.");
@@ -716,7 +738,7 @@ LogicalResult ConvertAtenOp<AtenSigmoidOp>::matchAndRewrite(
     ConversionPatternRewriter& rewriter) const {
   Location loc = op.getLoc();
   Value input = adaptor.self();
-  auto inputTy = input.getType().template cast<RankedTensorType>();
+  auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
   if (!inputTy) {
     return op.emitError(
         "Only RankedTensorType is supported in Aten Sigmoid op.");
@@ -738,7 +760,7 @@ LogicalResult ConvertAtenOp<AtenSiluOp>::matchAndRewrite(
     ConversionPatternRewriter& rewriter) const {
   Location loc = op.getLoc();
   Value input = adaptor.self();
-  auto inputTy = input.getType().template cast<RankedTensorType>();
+  auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
   if (!inputTy) {
     return op.emitError("Only RankedTensorType is supported in Aten SiLu op.");
   }
@@ -756,7 +778,7 @@ LogicalResult ConvertAtenOp<AtenGeluOp>::matchAndRewrite(
     ConversionPatternRewriter& rewriter) const {
   Location loc = op.getLoc();
   Value input = adaptor.self();
-  auto inputTy = input.getType().template cast<RankedTensorType>();
+  auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
   if (!inputTy) {
     return op.emitError("Only RankedTensorType is supported in Aten GELU op.");
   }
@@ -1255,7 +1277,7 @@ class ConvertAtenPoolingBaseOp : public OpConversionPattern<AtenOpT> {
       ConversionPatternRewriter& rewriter,
       Value input,
       ArrayRef<int32_t> transposeDims) const {
-    auto inputTy = input.getType().template cast<RankedTensorType>();
+    auto inputTy = input.getType().template dyn_cast<RankedTensorType>();
     auto inputElemTy = inputTy.getElementType();
     auto inputShape = inputTy.getShape();
     auto inputRank = inputTy.getRank();
@@ -1755,6 +1777,7 @@ class ConvertTorchToMhlo
     INSERT_ATENOP_PATTERN(AtenBroadcastToOp);
     INSERT_ATENOP_PATTERN(AtenSigmoidOp);
     INSERT_ATENOP_PATTERN(AtenReluOp);
+    INSERT_ATENOP_PATTERN(AtenRelu6Op);
     INSERT_ATENOP_PATTERN(AtenLeakyReluOp);
     INSERT_ATENOP_PATTERN(AtenSiluOp);
     INSERT_ATENOP_PATTERN(AtenGeluOp);
diff --git a/pytorch_blade/third_party/torch-mlir b/pytorch_blade/third_party/torch-mlir
index 1a4b37b0a3f..105bfad4e9f 160000
--- a/pytorch_blade/third_party/torch-mlir
+++ b/pytorch_blade/third_party/torch-mlir
@@ -1,1 +1,1 @@
-Subproject commit 1a4b37b0a3f197c1fbce1943df31b0688d2f3121
+Subproject commit 105bfad4e9fe5e6175f301ece45c2b8e53d326e6


Test patch:
diff --git a/pytorch_blade/tests/disc/ops/test_activation.py b/pytorch_blade/tests/disc/ops/test_activation.py
index e7a91670ce0..69cc644e4b7 100644
--- a/pytorch_blade/tests/disc/ops/test_activation.py
+++ b/pytorch_blade/tests/disc/ops/test_activation.py
@@ -62,14 +62,11 @@ def test_hardtanh(self):
         self._test_activation(torch.nn.Hardtanh(), torch.nn.functional.hardtanh, [([2, 4, 16, 16], torch.float)])
 
     @skipIfEnableTorchMlir()
-    #TODO(yancey1989): dependence torch dialect add glu op
     def test_glu(self):
         self._test_activation(torch.nn.GLU(), torch.nn.functional.glu, [([2, 4, 16, 16], torch.float)])
 
-   
-    @skipIfEnableTorchMlir()
-    #TODO(yancey1989): need a white list which is separate from old mhlo converter
     def test_hardswish(self):
+
         def _jit_pass_hardswish(graph):
             from_graph_str = """
             graph(%x):
@@ -94,10 +91,11 @@ def hard_swish(x: torch.Tensor) -> torch.Tensor:
             torch._C._jit_pass_dce(graph)
             torch._C._jit_pass_constant_propagation(graph)
 
-        config = Config.get_current_context_or_new()
-        config.customize_jit_passes = [_jit_pass_hardswish]
-        with config:
+        with Config.get_current_context_or_new() as cfg:
+            cfg.customize_jit_passes = [_jit_pass_hardswish]
             self._test_activation(None, torch.nn.functional.hardswish, [([2, 4, 16, 16], torch.float)])
+            self._test_activation(None, torch.nn.functional.hardswish, [([-1, -1, -1, -1], torch.float)])
+
 
 
 if __name__ == "__main__":
diff --git a/pytorch_blade/tests/disc/ops/test_binary_ops.py b/pytorch_blade/tests/disc/ops/test_binary_ops.py
index f04da1ddd14..589f58fc065 100644
--- a/pytorch_blade/tests/disc/ops/test_binary_ops.py
+++ b/pytorch_blade/tests/disc/ops/test_binary_ops.py
@@ -11,6 +11,7 @@
 
 import torch
 import unittest
+import pytest
 
 from tests.disc.testing_base import DiscTestCase, skipIfEnableTorchMlir, isTorchMlirEnable
 
diff --git a/pytorch_blade/tests/disc/ops/test_unary_ops.py b/pytorch_blade/tests/disc/ops/test_unary_ops.py
index 2bfbaf43491..4717ebead48 100644
--- a/pytorch_blade/tests/disc/ops/test_unary_ops.py
+++ b/pytorch_blade/tests/disc/ops/test_unary_ops.py
@@ -8,7 +8,7 @@
 # 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.
-
+import pytest
 import torch
 import unittest
 from unittest import skipIf