Add dropout op (#436)

Co-authored-by: dan <dan@nod-labs.com>

e2e_testing/torchscript/basic.py

@@ -738,3 +738,23 @@ class AddCDivModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: AddCDivModule())
 def AddCDivModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(1,3), tu.rand(1,3), tu.rand(1,3))
+
+
+# ==============================================================================
+
+class DropoutModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+    def forward(self, x):
+        return torch.dropout(x, 0.0, False)
+
+
+@register_test_case(module_factory=lambda: DropoutModule())
+def DropoutModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4))

include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td

@@ -2199,6 +2199,22 @@ def Torch_AtenIntTensorOp : Torch_Op<"aten.Int.Tensor", [
   let assemblyFormat = "$a attr-dict `:` type($a) `->` type($result)";
 }
 
+def Torch_AtenDropoutOp : Torch_Op<"aten.dropout", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::dropout : (Tensor, float, bool) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$input,
+    Torch_FloatType:$p,
+    Torch_BoolType:$train
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$input `,` $p `,` $train attr-dict `:` type($input) `,` type($p) `,` type($train) `->` type($result)";
+}
+
 def Torch_Aten__Contains__StrOp : Torch_Op<"aten.__contains__.str", [
     AllowsTypeRefinement,
     HasValueSemantics

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -1132,6 +1132,33 @@ public:
 };
 } // namespace
 
+namespace {
+class ConvertAtenDropoutOp : public OpConversionPattern<AtenDropoutOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(AtenDropoutOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
+      return failure();
+
+    bool train;
+    if (!matchPattern(op.train(), m_TorchConstantBool(&train)))
+      return rewriter.notifyMatchFailure(op,
+                                         "Expected train to be constant bool.");
+
+    if (train)
+      return failure();
+    auto resultType = getTypeConverter()
+                          ->convertType(op->getResult(0).getType())
+                          .cast<RankedTensorType>();
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType,
+                                                adaptor.input());
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 // See comments at in convertMmOp and the heading for this section for general
 // considerations. This function needs to be auto-generated.
@@ -3035,6 +3062,8 @@ public:
     patterns.add<ConvertAtenIntTensorOp>(typeConverter, context);
     target.addIllegalOp<PrimNumToTensorScalarOp>();
     patterns.add<ConvertPrimNumToTensorScalarOp>(typeConverter, context);
+    target.addIllegalOp<AtenDropoutOp>();
+    patterns.add<ConvertAtenDropoutOp>(typeConverter, context);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -231,7 +231,7 @@ public:
             AtenContiguousOp, AtenFill_ScalarOp, AtenDetachOp,
             AtenMaskedFill_ScalarOp, AtenCopy_Op, AtenIndexPut_Op, AtenCumsumOp,
             AtenLayerNormOp, AtenClampOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
-            AtenLog2Op, Aten_SoftmaxBackwardDataOp, AtenRsqrtOp,
+            AtenLog2Op, Aten_SoftmaxBackwardDataOp, AtenRsqrtOp, AtenDropoutOp,
             AtenTanhBackwardOp, Aten_LogSoftmaxBackwardDataOp, AtenAddIntOp>(
             op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -569,6 +569,7 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
         emit("aten::IntImplicit : (Tensor) -> (int)")
         emit("aten::tensor.float : (float, int?, Device?, bool) -> (Tensor)")
         emit("aten::Int.Tensor : (Tensor) -> (int)")
+        emit("aten::dropout : (Tensor, float, bool) -> (Tensor)")
 
         # Dict ops.
         emit("aten::__contains__.str : (Dict(str, t), str) -> (bool)", has_folder=True)