Repository: alibaba/BladeDISC (language: python)

Issue:
enhancement bert training op clustering
To utilize BladeDISC accelerating bert training, we need to cluster more op into disc, as the experiment, we found 244 `FusionGroups` in the first iteration with LTC, to reduce the clusters, we should:

- enhancement op coverage, #441 lists all ops that required, I just list some priority ones:
    - native_layer_norm
    - _softmax
    - _log_softmax
    - _log_softmax_backward_data
    - _nll_loss_backward
    - _nll_loss_forward
    - tanh_backward
    - native_dropout_backward
    - gelu_backward
- enhancement shape propagation
    - native_layer_norm
    - native_dropout
    - native_dropout_backward
    

[here](https://gist.github.com/Yancey1989/2ffbdea3c48baa0243e56d436eecfc98) list the 224 FusionGroups


Discussion / hints:

Fix patch:
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
index e946cbf0f9b..169b35dc13a 100644
--- a/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.cpp
+++ b/pytorch_blade/src/compiler/mlir/converters/torch_mlir_op_filter.cpp
@@ -55,6 +55,7 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "aten::flip",
       "aten::floor_divide",
       "aten::gelu",
+      "aten::gelu_backward",
       "aten::glu",
       "aten::hardtanh",
       "aten::index_select",
@@ -70,6 +71,7 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "aten::ne",
       "aten::neg",
       "aten::neg",
+      "aten::native_dropout",
       "aten::permute",
       "aten::relu",
       "aten::relu6",
diff --git a/pytorch_blade/src/ltc/disc_compiler/passes/register_disc_class.cpp b/pytorch_blade/src/ltc/disc_compiler/passes/register_disc_class.cpp
index f9307865ca7..a420a9bbcdd 100644
--- a/pytorch_blade/src/ltc/disc_compiler/passes/register_disc_class.cpp
+++ b/pytorch_blade/src/ltc/disc_compiler/passes/register_disc_class.cpp
@@ -38,6 +38,8 @@ std::string DiscCMD(
   std::stringstream ss;
   std::string logf = mlir_fname + ".log";
   std::string binary_path = CurrentLibLocation() + "/disc_compiler_main";
+  if (GRAPH_DEBUG_ENABLED)
+    binary_path = "TF_CPP_VMODULE=disc_compiler=1 " + binary_path;
   ss << binary_path << " " << mlir_fname << " " << out_fname << " > " << logf
      << " 2>&1 ";
   return ss.str();
@@ -62,13 +64,12 @@ std::tuple<std::string, std::string, std::string> MhloConversaion(
   return std::make_tuple(in_fname, input_dev_str, output_dev_str);
 }
 
-std::tuple<std::string, std::string> CallDiscCompiler(
+std::tuple<std::string, std::string, int> CallDiscCompiler(
     const std::string& mlir_fname) {
   std::string out_fname = mlir_fname + ".out";
   std::string cmd = DiscCMD(mlir_fname, out_fname);
-  TORCH_CHECK(
-      std::system(cmd.c_str()) == 0, "disc compile failed with cmd: " + cmd);
-  return {cmd, out_fname};
+  auto ret = std::system(cmd.c_str());
+  return {cmd, out_fname, ret};
 }
 
 const std::vector<torch::jit::Value*> ArrayToVector(
@@ -129,8 +130,8 @@ std::vector<c10::IValue> RegisterDiscClass(
       graph->nodes().end(),
       std::back_inserter(disc_nodes),
       [](torch::jit::Node* node) { return node->kind() == prim::FusionGroup; });
-
-  for (auto node : disc_nodes) {
+  for (size_t i = 0; i < disc_nodes.size(); ++i) {
+    auto node = disc_nodes[i];
     auto sub_graph = node->g(attr::Subgraph);
     auto state = std::make_shared<torch::blade::backends::EngineState>();
     std::vector<torch::blade::backends::EngineState::TensorType> inputs,
@@ -138,7 +139,12 @@ std::vector<c10::IValue> RegisterDiscClass(
     GRAPH_DUMP("Compile before mhlo conversion \n ", sub_graph);
     auto cvt_ret = MhloConversaion(sub_graph);
     auto ret = CallDiscCompiler(std::get<0>(cvt_ret) /*mlir file name*/);
+    auto ret_code = std::get<2>(ret);
     auto cmd = std::get<0>(ret);
+    if (ret_code != 0) {
+      GRAPH_DEBUG("disc compilation fallback to torchscript, cmd: " + cmd);
+      continue;
+    }
     auto output_fname = std::get<1>(ret);
     GRAPH_DEBUG(
         "disc compile fusionGroup ",
diff --git a/pytorch_blade/src/torch-mlir/BUILD b/pytorch_blade/src/torch-mlir/BUILD
index 817a362a1fa..b0e8dea4fc0 100644
--- a/pytorch_blade/src/torch-mlir/BUILD
+++ b/pytorch_blade/src/torch-mlir/BUILD
@@ -59,10 +59,11 @@ cc_library(
         "@llvm-project//mlir:Dialect",
         "@torch-mlir//:TorchMLIRTorchBackendTypeConversion",
         "@torch-mlir//:TorchMLIRTorchConversionDialect",
-	"@torch-mlir//:TorchMLIRTorchConversionPasses",
-	"@torch-mlir//:TorchMLIRTorchConversionPassesIncGen",
+	    "@torch-mlir//:TorchMLIRTorchConversionPasses",
+	    "@torch-mlir//:TorchMLIRTorchConversionPassesIncGen",
         "@torch-mlir//:TorchMLIRConversionPassesIncGen",
-	"@org_tensorflow//tensorflow/compiler/mlir/hlo:hlo",
+	    "@org_tensorflow//tensorflow/compiler/mlir/hlo:hlo",
+        "@org_tensorflow//tensorflow/compiler/mlir/disc:rng_uniform_custom_call",
     ]
 )
 
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 d0e5fca0376..93f594b082b 100644
--- a/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
+++ b/pytorch_blade/src/torch-mlir/lib/Conversion/TorchToMhlo/TorchToMhlo.cpp
@@ -39,6 +39,9 @@
 #include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.h"
 #include "torch-mlir/Dialect/TorchConversion/Transforms/BackendTypeConversion.h"
 
+#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
+#include "tensorflow/compiler/mlir/disc/IR/rng_uniform_custom_call_op.h"
+
 using namespace mlir;
 using namespace mlir::torch;
 using namespace mlir::torch::Torch;
@@ -792,6 +795,80 @@ LogicalResult ConvertAtenOp<AtenGeluOp>::matchAndRewrite(
   return success();
 }
 
+template <>
+LogicalResult ConvertAtenOp<AtenGeluBackwardOp>::matchAndRewrite(
+    AtenGeluBackwardOp op,
+    OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  Location loc = op.getLoc();
+  Value input = adaptor.self();
+  Value grad_output = adaptor.self();
+  Value half = chlo::getConstantLike(rewriter, loc, 0.5, input);
+  Value one = chlo::getConstantLike(rewriter, loc, 1.0, input);
+  Value m_2_sqrtpi = chlo::getConstantLike(
+      rewriter, loc, 1.12837916709551257389615890312154517, input);
+  Value m_sqrt1_2 = chlo::getConstantLike(
+      rewriter, loc, 0.707106781186547524400844362104849039, input);
+
+  Value alpha = rewriter.create<mhlo::MulOp>(loc, m_2_sqrtpi, m_sqrt1_2);
+  alpha = rewriter.create<mhlo::MulOp>(loc, alpha, half);
+  Value scratch = rewriter.create<chlo::ErfOp>(
+      loc, rewriter.create<mhlo::MulOp>(loc, input, m_sqrt1_2));
+
+  // dinput = exp(input * input * neg(half))
+  Value dinput = rewriter.create<mhlo::MulOp>(loc, input, input);
+  Value neg_half = rewriter.create<mhlo::NegOp>(loc, half);
+  dinput = rewriter.create<mhlo::MulOp>(loc, input, neg_half);
+
+  // result = grad_output * (half * (one + scratch) + input * dinput * alpha)
+  Value half_one = rewriter.create<mhlo::AddOp>(loc, one, scratch);
+  half_one = rewriter.create<mhlo::MulOp>(loc, half, half_one);
+
+  Value input_alpha = rewriter.create<mhlo::MulOp>(loc, input, dinput);
+  input_alpha = rewriter.create<mhlo::MulOp>(loc, input_alpha, alpha);
+
+  Value result = rewriter.create<mhlo::AddOp>(loc, half_one, input_alpha);
+  result = rewriter.create<mhlo::MulOp>(loc, grad_output, result);
+  rewriter.replaceOp(op, {result});
+  return success();
+}
+
+// Convert AtenEmptyMemoryFormatOp to const ones Tensor
+// please note: The AtenEmptyMemoryFormatOp is not according to value sematic,
+// and AtenEmptyMemoryFormatOp is a side-effected, to elimiate this op after
+// Canonicalizer pass, we lower this op to a constant ones op, this is a tracy
+// way.
+template <>
+LogicalResult ConvertAtenOp<AtenEmptyMemoryFormatOp>::matchAndRewrite(
+    AtenEmptyMemoryFormatOp op,
+    OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  auto outType = getTypeConverter()
+                     ->convertType(op.getType())
+                     .template dyn_cast<TensorType>();
+  auto loc = op.getLoc();
+  auto shape = adaptor.size();
+  SmallVector<Value, 4> dimSizes;
+  getListConstructElements(shape, dimSizes);
+  // BladeDISC use i32 as shape
+  std::for_each(dimSizes.begin(), dimSizes.end(), [&](Value& dSize) {
+    dSize = rewriter.create<ToI64Op>(loc, dSize).getResult();
+    // dimSize: cast i64 -> i32
+    dSize = rewriter.create<arith::TruncIOp>(loc, rewriter.getI32Type(), dSize);
+    return dSize;
+  });
+
+  auto mhloShape = rewriter.create<mlir::tensor::FromElementsOp>(loc, dimSizes);
+  auto constOp =
+      mhlo::getConstTensor<int32_t>(rewriter, op, {1.0}, {}).getValue();
+
+  auto result = rewriter.create<mhlo::DynamicBroadcastInDimOp>(
+      loc, outType, constOp, mhloShape, rewriter.getI64TensorAttr({}));
+
+  rewriter.replaceOp(op, {result});
+  return success();
+}
+
 using ReductionConvFunc = llvm::Optional<Value> (*)(
     PatternRewriter&,
     Operation*,
@@ -1731,8 +1808,13 @@ LogicalResult ConvertAtenOp<PrimNumToTensorScalarOp>::matchAndRewrite(
 
   Value constOp;
   double val;
-  if (!matchPattern(op.a(), m_TorchConstantFloat(&val)))
-    return op.emitError("input should be constant");
+  if (!matchPattern(op.a(), m_TorchConstantFloat(&val))) {
+    constOp = scalarToMhloTensor(
+        rewriter, op, adaptor.a(), outType.getElementType(), {});
+    rewriter.replaceOp(op, constOp);
+    return success();
+  }
+  // return op.emitError("input should be constant");
 
   if (failed(torchScalarToMhloTensor(
           rewriter,
@@ -1994,6 +2076,75 @@ LogicalResult ConvertAtenOp<AtenBroadcastToOp>::matchAndRewrite(
   return success();
 }
 
+StringAttr PackRandomUniformBackendConfig(
+    IntegerAttr seed,
+    IntegerAttr seed2,
+    PatternRewriter* rewriter) {
+  mhlo_disc::RngUniformBackendConfig config(
+      seed.getValue().getSExtValue(), seed2.getValue().getSExtValue());
+  std::string str;
+  llvm::raw_string_ostream ostream(str);
+  ostream << ::llvm::json::toJSON(config);
+  return rewriter->getStringAttr(ostream.str());
+}
+
+template <>
+LogicalResult ConvertAtenOp<ValsemVariantAtenUniformOp>::matchAndRewrite(
+    ValsemVariantAtenUniformOp op,
+    OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  auto inputTy = adaptor.self().getType().template cast<RankedTensorType>();
+  auto loc = op.getLoc();
+  if (!inputTy) {
+    op.emitError("input should be ranked tensor type.");
+  }
+  auto definingOp = op.self().getDefiningOp();
+  auto shape = definingOp->getOperand(0);
+  SmallVector<Value, 4> dimSizes;
+  getListConstructElements(shape, dimSizes);
+  std::for_each(dimSizes.begin(), dimSizes.end(), [&](Value& dSize) {
+    dSize = rewriter.create<ToI64Op>(loc, dSize).getResult();
+    // dimSize: cast i64 -> i32
+    dSize = rewriter.create<arith::TruncIOp>(loc, rewriter.getI32Type(), dSize);
+    return dSize;
+  });
+
+  auto mhloShape =
+      rewriter.create<tensor::FromElementsOp>(op.getLoc(), dimSizes);
+
+  double fromDoubleValue, toDoubleValue;
+  if (!matchPattern(op.from(), m_TorchConstantFloat(&fromDoubleValue))) {
+    op.emitError("operand #1 should be scalar");
+  }
+  if (!matchPattern(op.to(), m_TorchConstantFloat(&toDoubleValue))) {
+    op.emitError("operand #2 should be scalar");
+  }
+  Value fromTensor = rewriter.create<mhlo::ConstantOp>(
+      op.getLoc(),
+      rewriter.getFloatAttr(inputTy.getElementType(), fromDoubleValue));
+  Value toTensor = rewriter.create<mhlo::ConstantOp>(
+      op.getLoc(),
+      rewriter.getFloatAttr(inputTy.getElementType(), toDoubleValue));
+
+  auto cfg = PackRandomUniformBackendConfig(
+      rewriter.getIntegerAttr(rewriter.getI32Type(), 1),
+      rewriter.getIntegerAttr(rewriter.getI32Type(), 2),
+      &rewriter);
+  auto outType = getTypeConverter()
+                     ->convertType(op.getType())
+                     .template dyn_cast<TensorType>();
+
+  auto custom_call_op = rewriter.create<mhlo_disc::CustomCallOp>(
+      op.getLoc(),
+      TypeRange{outType},
+      ValueRange{fromTensor, toTensor, mhloShape},
+      rewriter.getStringAttr("rng_uniform"),
+      rewriter.getBoolAttr(false),
+      cfg);
+  rewriter.replaceOp(op, {custom_call_op.getResult(0)});
+  return success();
+}
+
 template <typename AtenOpT, typename MhloOpT>
 class ConvertAtenPoolingBaseOp : public OpConversionPattern<AtenOpT> {
  public:
@@ -2372,8 +2523,9 @@ class ConvertAtenFillScalarOp : public OpConversionPattern<AtenOpT> {
                        .template dyn_cast<TensorType>();
 
     if (!outType || !outType.hasStaticShape())
-      return op.emitError(
-          "Only Tensor types with static shapes are currently supported");
+      if (!outType)
+        return op.emitError(
+            "Only Tensor types with static shapes are currently supported");
 
     Type outElemTy = outType.getElementType();
     if (!outElemTy.isIntOrFloat()) {
@@ -2406,6 +2558,7 @@ class ConvertTorchToMhlo
     registry.insert<mhlo::MhloDialect>();
     registry.insert<tensor::TensorDialect>();
     registry.insert<arith::ArithmeticDialect>();
+    registry.insert<mhlo_disc::MhloDiscDialect>();
     TorchConversion::getBackendTypeConversionDependentDialects(registry);
   }
 
@@ -2415,6 +2568,7 @@ class ConvertTorchToMhlo
     target.addLegalDialect<
         chlo::ChloDialect,
         mhlo::MhloDialect,
+        mhlo_disc::MhloDiscDialect,
         tensor::TensorDialect,
         arith::ArithmeticDialect>();
 
@@ -2472,8 +2626,8 @@ class ConvertTorchToMhlo
         AtenGtScalarOp, mhlo::ComparisonDirection::GT);
     // INSERT_BINARY_COMPARE_PATTERN(AtenGeTensorOp,
     // mhlo::ComparisonDirection::GE);
-    // INSERT_BINARY_COMPARE_PATTERN(AtenGeScalarOp,
-    // mhlo::ComparisonDirection::GE);
+    INSERT_BINARY_COMPARE_PATTERN(
+        AtenGeScalarOp, mhlo::ComparisonDirection::GE);
     INSERT_BINARY_COMPARE_PATTERN(
         AtenEqTensorOp, mhlo::ComparisonDirection::EQ);
     INSERT_BINARY_COMPARE_PATTERN(
@@ -2586,8 +2740,10 @@ class ConvertTorchToMhlo
     INSERT_ATENOP_PATTERN(AtenFlipOp);
     INSERT_ATENOP_PATTERN(AtenIndexSelectOp);
     INSERT_ATENOP_PATTERN(AtenRollOp);
+    INSERT_ATENOP_PATTERN(ValsemVariantAtenUniformOp);
     INSERT_ATENOP_PATTERN(TensorStaticInfoCastOp);
-    // INSERT_ATENOP_PATTERN(AtenGeluBackwardOp);
+    INSERT_ATENOP_PATTERN(AtenGeluBackwardOp);
+    INSERT_ATENOP_PATTERN(AtenEmptyMemoryFormatOp);
 #undef INSERT_ATENOP_PATTERN
 
 #define INSERT_VIEW_OP_PATTERN(AtenOp) \
diff --git a/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp b/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
index ac8356c3921..c32b853d75d 100644
--- a/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
+++ b/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp
@@ -78,6 +78,45 @@ LogicalResult ConvertAtenOp<AtenHardtanhOp>::matchAndRewrite(
 
   return success();
 }
+
+template <>
+LogicalResult ConvertAtenOp<AtenNativeDropoutOp>::matchAndRewrite(
+    AtenNativeDropoutOp op,
+    OpAdaptor adaptor,
+    ConversionPatternRewriter& rewriter) const {
+  Location loc = op.getLoc();
+  Value input = op.input();
+  Value prob = op.p();
+  bool train = false;
+  if (!matchPattern(op.train(), m_TorchConstantBool(&train)))
+    return rewriter.notifyMatchFailure(op, "train must be a boolean constant");
+
+  BaseTensorType inputType = input.getType().cast<BaseTensorType>();
+  if (!train) {
+    // TODO(yancey.yx): supports inference mode
+    return op.emitError(
+        "native_dropout does not support argument train is false");
+  }
+  if (!inputType.hasDtype() || !inputType.getDtype().isa<mlir::FloatType>())
+    return rewriter.notifyMatchFailure(
+        op, "only support floating type input for training mode");
+  Value noneVal = rewriter.create<ConstantNoneOp>(loc);
+  Value floatOne =
+      rewriter.create<ConstantFloatOp>(loc, rewriter.getF64FloatAttr(1.0));
+  Value oneMinusP = rewriter.create<AtenSubFloatOp>(loc, floatOne, prob);
+  Value boolMask = rewriter.create<ValsemVariantAtenBernoulliFloatOp>(
+      loc, inputType, input, oneMinusP, /*generator=*/noneVal);
+  Value maskedInput =
+      rewriter.create<AtenMulTensorOp>(loc, inputType, boolMask, input);
+  Value output =
+      rewriter.create<AtenMulScalarOp>(loc, inputType, maskedInput, oneMinusP);
+  Value one =
+      rewriter.create<Torch::ConstantIntOp>(loc, rewriter.getI64IntegerAttr(1));
+  boolMask = rewriter.create<AtenGeScalarOp>(
+      loc, op.getResult(1).getType(), boolMask, one);
+  rewriter.replaceOp(op, {output, boolMask});
+  return success();
+}
 } // namespace
 
 namespace {
@@ -94,6 +133,9 @@ class DiscDecomposeComplexOpsPass
     patterns.add<ConvertAtenOp<AtenHardtanhOp>>(context);
     target.addIllegalOp<AtenHardtanhOp>();
 
+    patterns.add<ConvertAtenOp<AtenNativeDropoutOp>>(context);
+    target.addIllegalOp<AtenNativeDropoutOp>();
+
     if (failed(applyPartialConversion(
             getOperation(), target, std::move(patterns)))) {
       return signalPassFailure();
diff --git a/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp b/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp
index 71f5d1f92c8..b9e180d3616 100644
--- a/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp
+++ b/pytorch_blade/src/torch-mlir/lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp
@@ -30,6 +30,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 
+#include "tensorflow/compiler/mlir/disc/IR/hlo_disc_ops.h"
 #include "torch-mlir/Conversion/MhloPasses.h"
 #include "torch-mlir/Conversion/TorchToMhlo/TorchToMhlo.h"
 #include "torch-mlir/Conversion/TorchToSCF/TorchToSCF.h"
@@ -150,6 +151,7 @@ class VerifyMhloBackendContractPass
         arith::ArithmeticDialect,
         chlo::ChloDialect,
         mhlo::MhloDialect,
+        mhlo_disc::MhloDiscDialect,
         tensor::TensorDialect>();
     RewritePatternSet patterns(context);
     populateReturnOpTypeConversionPattern(patterns, converter);
diff --git a/tao_compiler/mlir/disc/BUILD b/tao_compiler/mlir/disc/BUILD
index f5da463b880..8cbb049de3b 100644
--- a/tao_compiler/mlir/disc/BUILD
+++ b/tao_compiler/mlir/disc/BUILD
@@ -1942,6 +1942,7 @@ cc_library(
         "@llvm-project//mlir:IR",
     ],
     alwayslink = 1,
+    visibility = ["//visibility:public"],
 )
 
 cc_library(


Test patch:
diff --git a/pytorch_blade/tests/__init__.py b/pytorch_blade/tests/__init__.py
index 617e233e8b2..86354c11164 100644
--- a/pytorch_blade/tests/__init__.py
+++ b/pytorch_blade/tests/__init__.py
@@ -8,3 +8,9 @@
 # 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.
+
+try:
+    import onnx
+except Exception:
+    # note: import onnx before torch_blade to avoid pb conflict
+    pass
diff --git a/pytorch_blade/tests/mhlo/elementwise.mlir b/pytorch_blade/tests/mhlo/elementwise.mlir
index 7d97fbf60e7..313523dc75f 100644
--- a/pytorch_blade/tests/mhlo/elementwise.mlir
+++ b/pytorch_blade/tests/mhlo/elementwise.mlir
@@ -120,3 +120,57 @@ func.func @torch.aten.ne.scalar.const(%arg0: !torch.vtensor<[?,?],si32>) -> !tor
   %0 = torch.aten.gt.Scalar %arg0, %float2.000000e00 : !torch.vtensor<[?,?],si32>, !torch.float -> !torch.vtensor<[?,?],i1>
   return %0 : !torch.vtensor<[?,?],i1>
 }
+
+// CHECK-LABEL:   func @torch.aten.gelu_backward(
+// CHECK-SAME:            %[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
+// CHECK:           %[[CST_0:.*]] = "chlo.constant_like"(%[[ARG1]]) {value = 5.000000e-01 : f32} : (tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK:           %[[CST_1:.*]] = "chlo.constant_like"(%[[ARG1]]) {value = 1.000000e+00 : f32} : (tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK:           %[[CST_2:.*]] = "chlo.constant_like"(%[[ARG1]]) {value = 1.12837923 : f32} : (tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK:           %[[CST_3:.*]] = "chlo.constant_like"(%[[ARG1]]) {value = 0.707106769 : f32} : (tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK:           %[[T0:.*]] = mhlo.multiply %[[CST_2]], %[[CST_3]] : tensor<?x?xf32>
+// CHECK:           %[[T1:.*]] = mhlo.multiply %[[T0]], %[[CST_0]] : tensor<?x?xf32>
+// CHECK:           %[[T2:.*]] = mhlo.multiply %[[ARG1]], %[[CST_3]] : tensor<?x?xf32>
+// CHECK:           %[[T3:.*]] = chlo.erf %[[T2]] : tensor<?x?xf32> -> tensor<?x?xf32>
+// CHECK:           %[[T4:.*]] = mhlo.negate %[[CST_0]] : tensor<?x?xf32>
+// CHECK:           %[[T5:.*]] = mhlo.multiply %[[ARG1]], %[[T4]] : tensor<?x?xf32>
+// CHECK:           %[[T6:.*]] = mhlo.add %[[CST_1]], %[[T3]] : tensor<?x?xf32>
+// CHECK:           %[[T7:.*]] = mhlo.multiply %[[CST_0]], %[[T6]] : tensor<?x?xf32>
+// CHECK:           %[[T8:.*]] = mhlo.multiply %[[ARG1]], %[[T5]] : tensor<?x?xf32>
+// CHECK:           %[[T9:.*]] = mhlo.multiply %[[T8]], %[[T1]] : tensor<?x?xf32>
+// CHECK:           %[[T10:.*]] = mhlo.add %[[T7]], %[[T9]] : tensor<?x?xf32>
+// CHECK:           %[[T11:.*]] = mhlo.multiply %[[ARG1]], %[[T10]] : tensor<?x?xf32>
+// CHECK:           return %[[T11]] : tensor<?x?xf32>
+func.func @torch.aten.gelu_backward(%arg0: !torch.vtensor<[?,?],f32>, %arg1: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[?,?],f32> {
+  %str_none = torch.constant.str "none"
+  %0 = torch.aten.gelu_backward %arg0, %arg1, %str_none: !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],f32>, !torch.str -> !torch.vtensor<[?,?],f32>
+  return %0 : !torch.vtensor<[?,?],f32>
+}
+
+// CHECK-LABEL:   func @torch.aten.native_dropout.train(
+// CHECK-SAME:            %[[ARG0:.*]]: tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xi1>) {
+// CHECK:           %[[CST_0:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f32>
+// CHECK:           %[[CST_1:.*]] = mhlo.constant dense<0.000000e+00> : tensor<f64>
+// CHECK:           %[[CST_2:.*]] =  mhlo.constant dense<1.000000e+00> : tensor<f32>
+// CHECK:           %[[SHAPE_0:.*]] = arith.constant 1 : index
+// CHECK:           %[[SHAPE_1:.*]] = arith.constant 0 : index
+// CHECK:           %[[DIM_0:.*]] = tensor.dim %[[ARG0]], %[[SHAPE_1]] : tensor<?x?xf32>
+// CHECK:           %[[DIM_0_I64:.*]] = arith.index_cast %[[DIM_0]] : index to i64
+// CHECK:           %[[DIM_1:.*]] = tensor.dim %[[ARG0]], %[[SHAPE_0]] : tensor<?x?xf32>
+// CHECK:           %[[DIM_1_I64:.*]] = arith.index_cast %[[DIM_1]] : index to i64
+// CHECK:           %[[DIM_0_I32:.*]] = arith.trunci %[[DIM_0_I64]] : i64 to i32
+// CHECK:           %[[DIM_1_I32:.*]] = arith.trunci %[[DIM_1_I64]] : i64 to i32
+// CHECK:           %[[T0:.*]] = tensor.from_elements %[[DIM_0_I32]], %[[DIM_1_I32]] : tensor<2xi32>
+// CHECK:           %[[T1:.*]] = "mhlo_disc.custom_call"(%[[CST_0]], %[[CST_2]], %[[T0]]) {backend_config = "{\22seed\22:1,\22seed2\22:2}", call_target_name = "rng_uniform", has_side_effect = false} : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<?x?xf32>
+// CHECK:           %[[T2:.*]] = mhlo.convert(%[[T1]]) : (tensor<?x?xf32>) -> tensor<?x?xf64>
+// CHECK:           %[[T3:.*]] = chlo.broadcast_compare %[[T2]], %[[CST_1]] {compare_type = #mhlo<comparison_type NOTYPE>, comparison_direction = #mhlo<comparison_direction LT>} : (tensor<?x?xf64>, tensor<f64>) -> tensor<?x?xi1>
+// CHECK:           %[[T4:.*]] = mhlo.convert(%[[T3]]) : (tensor<?x?xi1>) -> tensor<?x?xf32>
+// CHECK:           %[[T5:.*]] = chlo.broadcast_multiply %[[T4]], %[[ARG0]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+// CHECK:           %[[T6:.*]] = chlo.broadcast_multiply %[[T5]], %[[CST_0]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
+// CHECK:           %[[T7:.*]] = chlo.broadcast_compare %[[T4]], %[[CST_2]] {compare_type = #mhlo<comparison_type NOTYPE>, comparison_direction = #mhlo<comparison_direction GE>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
+// CHECK:           return %[[T6]], %[[T7]] : tensor<?x?xf32>, tensor<?x?xi1>
+func.func @torch.aten.native_dropout.train(%arg0: !torch.vtensor<[?,?],f32>) -> (!torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],i1>) {
+  %float1 = torch.constant.float 1.000000e+00
+  %bool_true = torch.constant.bool true
+  %result0, %result1 = torch.aten.native_dropout %arg0, %float1, %bool_true: !torch.vtensor<[?,?],f32>, !torch.float, !torch.bool -> !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],i1>
+  return %result0, %result1 : !torch.vtensor<[?,?],f32>, !torch.vtensor<[?,?],i1>
+}
\ No newline at end of file