mirror of https://github.com/llvm/torch-mlir
parent
03fdf56f21
commit
539511c19b
|
@ -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))
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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))))
|
||||
|
|
|
@ -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]);
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue