Add lowering for aten.to.device (#1107)

pull/1176/head
gpetters94 2022-08-10 19:24:02 -04:00 committed by GitHub
parent b8d51a74d9
commit 79b9cf9468
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 82 additions and 1 deletions

View File

@ -53,6 +53,7 @@ TOSA_PASS_SET = {
"TModuleRank1_basic",
"TModuleRank0_basic",
"ElementwiseToDtypeIdentityModule_basic",
"AtenToDeviceModule_basic",
"View1DFoldModule_basic",
"UnsafeView1DFoldModule_basic",
"SqueezeDimModule_static",

View File

@ -5880,6 +5880,33 @@ def Torch_AtenToPrimDeviceOp : Torch_Op<"aten.to.prim_Device", [
}];
}
def Torch_AtenToDeviceOp : Torch_Op<"aten.to.device", [
AllowsTypeRefinement,
ReadOnly
]> {
let summary = "Generated op for `aten::to.device : (Tensor, Device, int, bool, bool, int?) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self,
Torch_DeviceType:$device,
Torch_IntType:$dtype,
Torch_BoolType:$non_blocking,
Torch_BoolType:$copy,
AnyTorchOptionalIntType:$memory_format
);
let results = (outs
AnyTorchTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenToDeviceOp::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 6, 1);
}
void AtenToDeviceOp::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 6, 1);
}
}];
}
def Torch_AtenTypeAsOp : Torch_Op<"aten.type_as", [
AllowsTypeRefinement,
HasValueSemantics,

View File

@ -1996,6 +1996,25 @@ public:
};
} // namespace
namespace {
// Decompose `aten.to.device` op into `aten.to.dtype` op.
class DecomposeAtenToDeviceOp : public OpRewritePattern<AtenToDeviceOp> {
public:
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(AtenToDeviceOp op,
PatternRewriter &rewriter) const override {
// Device information isn't relevant to torch-mlir, so we can drop that info
// here.
rewriter.replaceOpWithNewOp<AtenToDtypeOp>(op, op.getType(), op.self(),
op.dtype(), op.non_blocking(),
op.copy(), op.memory_format());
return success();
}
};
} // namespace
namespace {
// Decompose `aten.adaptive_avg_pool2d` op into `aten.avg_pool2d` op.
//
@ -2586,6 +2605,8 @@ class DecomposeComplexOpsPass
patterns.add<DecomposeAtenPadOp>(context);
patterns.add<DecomposeAtenToDtypeLayoutOp>(context);
target.addIllegalOp<AtenToDtypeLayoutOp>();
patterns.add<DecomposeAtenToDeviceOp>(context);
target.addIllegalOp<AtenToDeviceOp>();
patterns.add<DecomposeAtenAdaptiveAvgPool2dOp>(context);
target.addIllegalOp<AtenAdaptiveAvgPool2dOp>();
patterns.add<DecomposeAtenClampMinOp>(context);

View File

@ -38,7 +38,7 @@ static bool isViewLikeOp(Operation *op) {
AtenSqueezeDimOp, AtenSqueezeOp, AtenTOp, AtenToDtypeOp,
AtenTransposeIntOp, AtenUnsqueezeOp, AtenViewOp,
TensorStaticInfoCastOp, AtenToDtypeLayoutOp, AtenNumpyTOp,
AtenNarrowOp>(op);
AtenNarrowOp, AtenToDeviceOp>(op);
}
namespace {

View File

@ -1024,6 +1024,11 @@ void TypeAnalysis::visitOperation(Operation *op,
return;
}
if (auto toDtype = dyn_cast<AtenToDeviceOp>(op)) {
visitAtenToDtypeLikeOp<AtenToDeviceOp>(toDtype, operands);
return;
}
if (auto toOther = dyn_cast<AtenToOtherOp>(op)) {
visitTypeConversionOp<AtenToOtherOp>(toOther, operands);
return;

View File

@ -5448,6 +5448,10 @@ module {
func.func @"__torch_mlir_shape_fn.aten.to.dtype_layout"(%arg0: !torch.list<int>, %arg1: !torch.optional<int>, %arg2: !torch.optional<int>, %arg3: !torch.optional<Device>, %arg4: !torch.optional<bool>, %arg5: !torch.bool, %arg6: !torch.bool, %arg7: !torch.optional<int>) -> !torch.list<int> {
return %arg0 : !torch.list<int>
}
func.func @"__torch_mlir_shape_fn.aten.to.device"(%arg0: !torch.list<int>, %arg1: !torch.Device, %arg2: !torch.int, %arg3: !torch.bool, %arg4: !torch.bool, %arg5: !torch.optional<int>) -> !torch.list<int> {
%0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>
return %0 : !torch.list<int>
}
func.func @"__torch_mlir_shape_fn.aten.to.other"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.bool, %arg3: !torch.bool, %arg4: !torch.optional<int>) -> !torch.list<int> {
%0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>
return %0 : !torch.list<int>

View File

@ -427,6 +427,9 @@ def atentodtype(self: List[int], dtype: int, non_blocking: bool = False, c
def atentodtype_layout(self: List[int], dtype: Optional[int] = None, layout: Optional[int] = None, device: Optional[device] = None, pin_memory: Optional[bool] = None, non_blocking: bool = False, copy: bool = False, memory_format: Optional[int] = None) -> List[int]:
return self
def atentodevice(self: List[int], device: device, dtype: int, non_blocking: bool = False, copy: bool = False, memory_format: Optional[int] = None) -> List[int]:
return upstream_shape_functions.unary(self)
def atentoother(self: List[int], other: List[int], non_blocking: bool = False, copy: bool = False, memory_format: Optional[int] = None) -> List[int]:
return upstream_shape_functions.unary(self)

View File

@ -456,6 +456,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
emit("aten::to.dtype_layout : (Tensor, int?, int?, Device?, bool?, bool, bool, int?) -> (Tensor)", has_folder=True)
emit("aten::to.other : (Tensor, Tensor, bool, bool, int?) -> (Tensor)")
emit("aten::to.prim_Device : (Tensor, Device?, int?, bool, bool) -> (Tensor)")
emit("aten::to.device : (Tensor, Device, int, bool, bool, int?) -> (Tensor)")
emit("aten::type_as : (Tensor, Tensor) -> (Tensor)")
emit("aten::view : (Tensor, int[]) -> (Tensor)", has_folder=True)
emit("aten::_unsafe_view : (Tensor, int[]) -> (Tensor)")

View File

@ -2627,3 +2627,22 @@ def Aten_EmbeddingBagExample_basic(module, tu: TestUtils):
indices = torch.LongTensor([0, 1, 2, 2, 0, 2, 1, 3, 20, 50, 99, 2, 4, 5, 6, 7, 34, 54])
offsets = torch.LongTensor([0, 3, 5, 7, 9, 10, 15])
module.forward(weight, indices, offsets)
# ==============================================================================
class AtenToDeviceModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1 , -1], torch.float32, True),
])
def forward(self, val):
return torch.ops.aten.to(val, device='cpu', dtype=torch.float, non_blocking=False)
@register_test_case(module_factory=lambda: AtenToDeviceModule())
def AtenToDeviceModule_basic(module, tu: TestUtils):
module.forward(torch.randn(2, 4))