Add pytorch interface to ATen Dialect (#30)
This patch adds a pytorch interface to npcomp. This interface is modeled
after pytorch_xla and exposes the MLIR-based flow as a virtual device (similar
to a gpu device or the xla backend). Usage is intended to be something like:
dev = torch_mlir.mlir_device()
t0 = torch.randn((4,4), device=dev)
t1 = torch.randn((4,4), device=dev)
t2 = t0 + t1
t2_mlir = torch_mlir.get_mlir( t2 )
t2_cpu = t2.to('cpu')
In this case t2_cpu would contain the result of the computation, and t2_mlir
contains the mlir description of the computation. Note that this also
properly returns backward paths synthesized by pytorch. There are several
parts of this:
1) A tensor type (implemented by tensor.* and tensor_impl.*)
2) The device modeling (aten_mlir_bridge.*, aten_mlir_device.*, aten_mlir_type*)
3) a temporary IR (implemented by ir.cpp)
There is also a reference lowering directly from the ATen dialect to C
function calls consisting of two parts:
1) The driver that uses the IR to generate MLIR, run Passes and compile the
result using mlir::ExecutionEngine (implemented by jit.cpp and
mlir_gen.cpp)
2) A runtime library implemented by lib/aten_ops.cpp. Most of the operations
are implemented by callbacks into the torch C++ libraries.
Some aspects of this are known to be less than optimal, in particular:
1) There's some function definitions that don't live in the file corresponding
to their declaration.
2) More aspects of this (e.g. the IR) seem like they should be automatically
generated.
3) It's unclear to me how much of the 'IR' is actually necessary, or whether
MLIR could be created on the fly.
Note that this code is licensed in a way similar to pytorch, with the
intention that eventually (when npcomp reaches some maturity) it should be
pushed there. (see frontends/pytorch/LICENSE) The code is also structured
much closer to the pytorch coding style than the LLVM coding style.
2020-08-22 02:22:47 +08:00
|
|
|
# -*- Python -*-
|
|
|
|
# This file is licensed under a pytorch-style license
|
|
|
|
# See frontends/pytorch/LICENSE for license information.
|
|
|
|
|
|
|
|
import os
|
|
|
|
import platform
|
|
|
|
import re
|
|
|
|
import subprocess
|
|
|
|
import tempfile
|
|
|
|
|
|
|
|
import lit.formats
|
|
|
|
import lit.util
|
|
|
|
|
|
|
|
from lit.llvm import llvm_config
|
|
|
|
from lit.llvm.subst import ToolSubst
|
|
|
|
from lit.llvm.subst import FindTool
|
|
|
|
|
|
|
|
# Configuration file for the 'lit' test runner.
|
|
|
|
|
|
|
|
# name: The name of this test suite.
|
|
|
|
config.name = 'FRONTENDS_PYTORCH'
|
|
|
|
|
|
|
|
config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
|
2020-10-09 09:29:59 +08:00
|
|
|
config.environment['PYTHONPATH'] = os.path.join(config.npcomp_obj_root, "python")
|
Add pytorch interface to ATen Dialect (#30)
This patch adds a pytorch interface to npcomp. This interface is modeled
after pytorch_xla and exposes the MLIR-based flow as a virtual device (similar
to a gpu device or the xla backend). Usage is intended to be something like:
dev = torch_mlir.mlir_device()
t0 = torch.randn((4,4), device=dev)
t1 = torch.randn((4,4), device=dev)
t2 = t0 + t1
t2_mlir = torch_mlir.get_mlir( t2 )
t2_cpu = t2.to('cpu')
In this case t2_cpu would contain the result of the computation, and t2_mlir
contains the mlir description of the computation. Note that this also
properly returns backward paths synthesized by pytorch. There are several
parts of this:
1) A tensor type (implemented by tensor.* and tensor_impl.*)
2) The device modeling (aten_mlir_bridge.*, aten_mlir_device.*, aten_mlir_type*)
3) a temporary IR (implemented by ir.cpp)
There is also a reference lowering directly from the ATen dialect to C
function calls consisting of two parts:
1) The driver that uses the IR to generate MLIR, run Passes and compile the
result using mlir::ExecutionEngine (implemented by jit.cpp and
mlir_gen.cpp)
2) A runtime library implemented by lib/aten_ops.cpp. Most of the operations
are implemented by callbacks into the torch C++ libraries.
Some aspects of this are known to be less than optimal, in particular:
1) There's some function definitions that don't live in the file corresponding
to their declaration.
2) More aspects of this (e.g. the IR) seem like they should be automatically
generated.
3) It's unclear to me how much of the 'IR' is actually necessary, or whether
MLIR could be created on the fly.
Note that this code is licensed in a way similar to pytorch, with the
intention that eventually (when npcomp reaches some maturity) it should be
pushed there. (see frontends/pytorch/LICENSE) The code is also structured
much closer to the pytorch coding style than the LLVM coding style.
2020-08-22 02:22:47 +08:00
|
|
|
if 'TEST_SRC_PATH' in os.environ:
|
|
|
|
config.environment['TEST_SRC_PATH'] = os.environ['TEST_SRC_PATH']
|
|
|
|
|
|
|
|
# path to our python operation library
|
|
|
|
config.environment['TEST_BUILD_PATH'] = os.path.join(config.npcomp_obj_root)
|
|
|
|
|
|
|
|
# suffixes: A list of file extensions to treat as test files.
|
|
|
|
config.suffixes = ['.py']
|
|
|
|
|
|
|
|
# test_source_root: The root path where tests are located.
|
|
|
|
config.test_source_root = os.path.dirname(__file__)
|
|
|
|
|
|
|
|
# test_exec_root: The root path where tests should be run.
|
|
|
|
config.test_exec_root = os.path.join(config.npcomp_obj_root, 'test')
|
|
|
|
|
|
|
|
config.substitutions.append(('%PATH%', config.environment['PATH']))
|
|
|
|
config.substitutions.append(('%shlibext', config.llvm_shlib_ext))
|
|
|
|
|
|
|
|
llvm_config.with_system_environment(
|
|
|
|
['HOME', 'INCLUDE', 'LIB', 'TMP', 'TEMP'])
|
|
|
|
|
|
|
|
llvm_config.use_default_substitutions()
|
|
|
|
|
|
|
|
# excludes: A list of directories to exclude from the testsuite. The 'Inputs'
|
|
|
|
# subdirectories contain auxiliary inputs for various tests in their parent
|
|
|
|
# directories.
|
|
|
|
config.excludes = ['lit.cfg.py', 'Inputs', 'Examples', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt']
|
|
|
|
|
|
|
|
# test_source_root: The root path where tests are located.
|
|
|
|
config.test_source_root = os.path.dirname(__file__)
|
|
|
|
|
|
|
|
# test_exec_root: The root path where tests should be run.
|
|
|
|
config.test_exec_root = os.path.join(config.npcomp_obj_root, 'test')
|
|
|
|
config.npcomp_tools_dir = os.path.join(config.npcomp_obj_root, 'bin')
|
|
|
|
|
|
|
|
# Tweak the PATH to include the tools dir.
|
|
|
|
llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
|
|
|
|
|
|
|
|
tool_dirs = [config.npcomp_tools_dir, config.llvm_tools_dir]
|
|
|
|
tools = [
|
|
|
|
]
|
|
|
|
|
|
|
|
llvm_config.add_tool_substitutions(tools, tool_dirs)
|