llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/frontends/pytorch/test/node_import/add3.py

27 lines
1.4 KiB
Python
Raw Normal View History

Add TorchScript graph importer. * Does not handle all features yet but should conservatively fail on unsupported things. * Location tracking is still somewhat mismatched between what TorchScript and MLIR do. Likely need a better heuristic for tracking locations from defs for nodes that do not carry location. * Sets the ground-work for a specialized/generic split but only implements the generic side. * Had some evidence that this requires a recent bump of PT nightly (within the last month) to pick up pybind11 2.6, which includes some cross-module symbol fixes (vs the previously sync'd version). No source changes, but older versions fail to cast function types at runtime.
2020-11-21 09:03:23 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
import torch
import torch_mlir
# RUN: %PYTHON %s | npcomp-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
# Verify without debug info.
Add support for prim::GetAttr/SetAttr/CallMethod/If This required some invasive surgery to graph_importer.h/cpp, specifically moving most of it into node_importer.h/cpp and relayering it. The abstraction that it had didn't work well in the recursive setting that happens with prim::If. The key observation is that torch::jit::Graph doesn't really correspond directly to anything on the MLIR side. It's a weird combination of a context, builder, and function and just holds a `torch::jit::Block`. It is `torch::jit::Node` and `torch::jit::Block` which form the recursive structure analogous to MLIR's operation/region/block. So node_importer.h/cpp makes sense as a core building block. As part of doing this, I did venture a bit into the AcapController code, and realize now that there is functionality duplicated there with the ivalue importer. Will refactor that soon.
2021-02-02 09:59:42 +08:00
# CHECK-LABEL: func @add3
Add TorchScript graph importer. * Does not handle all features yet but should conservatively fail on unsupported things. * Location tracking is still somewhat mismatched between what TorchScript and MLIR do. Likely need a better heuristic for tracking locations from defs for nodes that do not carry location. * Sets the ground-work for a specialized/generic split but only implements the generic side. * Had some evidence that this requires a recent bump of PT nightly (within the last month) to pick up pybind11 2.6, which includes some cross-module symbol fixes (vs the previously sync'd version). No source changes, but older versions fail to cast function types at runtime.
2020-11-21 09:03:23 +08:00
# CHECK-SAME: (%arg0: !numpy.ndarray<*:!numpy.any_dtype>, %arg1: !numpy.ndarray<*:!numpy.any_dtype>, %arg2: !numpy.ndarray<*:!numpy.any_dtype>) -> !numpy.ndarray<*:!numpy.any_dtype> {
# CHECK: %[[C1:.*]] = constant 1 : i64
# CHECK: %[[A0:.*]] = torch.kernel_call "aten::add" %arg0, %arg1, %[[C1]] : (!numpy.ndarray<*:!numpy.any_dtype>, !numpy.ndarray<*:!numpy.any_dtype>, i64) -> !numpy.ndarray<*:!numpy.any_dtype> {sigArgTypes = ["Tensor", "Tensor", "Scalar"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor"]}
# CHECK: %[[A1:.*]] = torch.kernel_call "aten::add" %[[A0]], %arg2, %[[C1]] : (!numpy.ndarray<*:!numpy.any_dtype>, !numpy.ndarray<*:!numpy.any_dtype>, i64) -> !numpy.ndarray<*:!numpy.any_dtype> {sigArgTypes = ["Tensor", "Tensor", "Scalar"], sigIsMutable = false, sigIsVararg = false, sigIsVarret = false, sigRetTypes = ["Tensor"]}
# CHECK: return %[[A1]] : !numpy.ndarray<*:!numpy.any_dtype>
Add TorchScript import tests missed in previous change.
2020-11-24 06:41:30 +08:00
@mb.import_function
@torch.jit.script
def add3(t0, t1, t2):
return t0 + t1 + t2
assert isinstance(add3, torch.jit.ScriptFunction)
Add TorchScript graph importer. * Does not handle all features yet but should conservatively fail on unsupported things. * Location tracking is still somewhat mismatched between what TorchScript and MLIR do. Likely need a better heuristic for tracking locations from defs for nodes that do not carry location. * Sets the ground-work for a specialized/generic split but only implements the generic side. * Had some evidence that this requires a recent bump of PT nightly (within the last month) to pick up pybind11 2.6, which includes some cross-module symbol fixes (vs the previously sync'd version). No source changes, but older versions fail to cast function types at runtime.
2020-11-21 09:03:23 +08:00
mb.module.operation.print()
print()