mirror of https://github.com/llvm/torch-mlir
Add lowering for aten.to.device (#1107)
parent
b8d51a74d9
commit
79b9cf9468
|
@ -53,6 +53,7 @@ TOSA_PASS_SET = {
|
|||
"TModuleRank1_basic",
|
||||
"TModuleRank0_basic",
|
||||
"ElementwiseToDtypeIdentityModule_basic",
|
||||
"AtenToDeviceModule_basic",
|
||||
"View1DFoldModule_basic",
|
||||
"UnsafeView1DFoldModule_basic",
|
||||
"SqueezeDimModule_static",
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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 {
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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>
|
||||
|
|
|
@ -427,6 +427,9 @@ def aten〇to〇dtype(self: List[int], dtype: int, non_blocking: bool = False, c
|
|||
def aten〇to〇dtype_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 aten〇to〇device(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 aten〇to〇other(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)
|
||||
|
||||
|
|
|
@ -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)")
|
||||
|
|
|
@ -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))
|
Loading…
Reference in New Issue