llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Implement Python relayout from #311 

Fixes https://github.com/llvm/mlir-npcomp/issues/311

The key change is that TorchPlugin is folded into
`torch_mlir.dialects.torch.importer.jit_ir` (it imports the PyTorch
JIT's IR, so that's a good, scoped name for it).
The CMake option `-DTORCH_MLIR_ENABLE_JIT_IR_IMPORTER=OFF` disables it,
which allows building without a PyTorch native dependency.
pull/314/head
Sean Silva 2021-09-20 20:55:36 +00:00
parent ecc334123c
commit 6d8e7f1bb1
94 changed files with 270 additions and 559 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CMakeLists.txt
View File

@ -216,7 +216,6 @@ add_dependencies(check-npcomp-all
check-npcomp
check-npcomp-python
check-torch-mlir
check-torch-mlir-plugin
)
add_subdirectory(include/npcomp)

9
external/torch-mlir/CMakeLists.txt vendored
View File

@ -6,6 +6,7 @@ message(FATAL_ERROR
endif()
option(MLIR_ENABLE_BINDINGS_PYTHON "Enables MLIR Python Bindings" OFF)
option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
set(TORCH_MLIR_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
set(TORCH_MLIR_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
@ -66,16 +67,12 @@ add_subdirectory(lib)
add_subdirectory(tools)
if(MLIR_ENABLE_BINDINGS_PYTHON)
# If parent projects want to configure where to place the python packages,
# respect that.
if(NOT TORCH_MLIR_PYTHON_PACKAGES_DIR)
set(TORCH_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
endif()
add_subdirectory(TorchPlugin)
add_subdirectory(python)
endif()
add_subdirectory(test)
add_custom_target(check-torch-mlir-all)
add_dependencies(check-torch-mlir-all
check-torch-mlir
check-torch-mlir-plugin
)

0
external/torch-mlir/TorchPlugin/LICENSE → external/torch-mlir/LICENSE.pytorch vendored
View File

41
external/torch-mlir/TorchPlugin/CMakeLists.txt vendored
View File

@ -1,41 +0,0 @@
#-------------------------------------------------------------------------------
# Sub project setup
#-------------------------------------------------------------------------------
cmake_minimum_required(VERSION 3.13.4)
if(POLICY CMP0068)
cmake_policy(SET CMP0068 NEW)
set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
endif()
if(POLICY CMP0075)
cmake_policy(SET CMP0075 NEW)
endif()
if(POLICY CMP0077)
cmake_policy(SET CMP0077 NEW)
endif()
project(torch_mlir_plugin LANGUAGES CXX C)
set(CMAKE_C_STANDARD 11)
set(CMAKE_CXX_STANDARD 14)
#-------------------------------------------------------------------------------
# Setup PyTorch
#-------------------------------------------------------------------------------
list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
include(TorchMLIRPyTorch)
TorchMLIRProbeForPyTorchInstall()
find_package(Torch 1.8 REQUIRED)
TorchMLIRConfigurePyTorch()
#-------------------------------------------------------------------------------
# Subdirectories
#-------------------------------------------------------------------------------
add_subdirectory(csrc)
add_subdirectory(python)
add_subdirectory(test)

29
external/torch-mlir/TorchPlugin/csrc/init_python_bindings.cpp vendored
View File

@ -1,29 +0,0 @@
//===- init_python_bindings.cpp ---------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
//
//===----------------------------------------------------------------------===//
// This is the top-level entry point for the MLIR <-> PyTorch bridge.
#include "init_python_bindings.h"
#include <string>
namespace py = pybind11;
namespace torch_mlir {
namespace {
void InitModuleBindings(py::module &m) {}
} // namespace
void InitBindings(py::module &m) {
InitModuleBindings(m);
InitBuilderBindings(m);
}
} // namespace torch_mlir
PYBIND11_MODULE(_torch_mlir, m) { torch_mlir::InitBindings(m); }

23
external/torch-mlir/TorchPlugin/csrc/init_python_bindings.h vendored
View File

@ -1,23 +0,0 @@
//===- init_python_bindings.h -----------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
//
//===----------------------------------------------------------------------===//
#ifndef INIT_PYTHON_BINDINGS_H
#define INIT_PYTHON_BINDINGS_H
#include "pybind.h"
namespace torch_mlir {
// Perform top-level initialization for the module.
void InitBindings(pybind11::module &m);
// Adds bindings related to building modules.
void InitBuilderBindings(pybind11::module &m);
} // namespace torch_mlir
#endif

27
external/torch-mlir/TorchPlugin/python/CMakeLists.txt vendored
View File

@ -1,27 +0,0 @@
## Declare the sources of the Python module.
declare_mlir_python_sources(TorchMLIRPluginPythonSources)
declare_mlir_python_sources(TorchMLIRPluginPythonSources.Core
ADD_TO_PARENT TorchMLIRPluginPythonSources
SOURCES_GLOB
torch_mlir/*.py
torch_mlir_utils/codegen/*.py
)
set(_source_components
TorchMLIRPluginPythonSources
)
add_mlir_python_modules(TorchMLIRPluginPythonModules
ROOT_PREFIX "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir"
INSTALL_PREFIX "python_packages/torch_mlir"
DECLARED_SOURCES ${_source_components}
COMMON_CAPI_LINK_LIBS
TorchMLIRAggregateCAPI
)
add_dependencies(TorchMLIRPluginPythonModules
TorchMLIRTorchPlugin
TorchMLIRPythonModules
)

49
external/torch-mlir/TorchPlugin/setup.py vendored
View File

@ -1,49 +0,0 @@
# Build/install the npcomp-torch package.
# This uses PyTorch's setuptools support and requires an existing installation
# of npcomp-core in order to access its headers/libraries.
from pathlib import Path
from setuptools import find_packages, setup, Extension
from torch.utils import cpp_extension
try:
from npcomp import build as npcomp_build
except ModuleNotFoundError:
raise ModuleNotFoundError(
f"Could not import npcomp.build "
f"(do you have the npcomp-core package installed)")
# Get our sources.
this_dir = Path(__file__).parent
extension_sources = [str(p) for p in this_dir.joinpath("csrc").rglob("*.cpp")]
# Npcomp bits.
include_dirs = npcomp_build.get_include_dirs()
lib_dirs = npcomp_build.get_lib_dirs()
npcomp_libs = [npcomp_build.get_capi_link_library_name()]
# TODO: Add a way to customize MLIR_PYTHON_PACKAGE_PREFIX
compile_args = []
setup(
name="npcomp-torch",
ext_modules=[
cpp_extension.CppExtension(
name="_torch_mlir",
sources=extension_sources,
include_dirs=include_dirs,
library_dirs=lib_dirs,
libraries=npcomp_libs,
extra_compile_args=compile_args),
],
cmdclass={
"build_ext": cpp_extension.BuildExtension
},
package_dir={
"": "./python",
},
packages=find_packages("./python", include=[
"torch_mlir",
"torch_mlir.*",
]),
)

21
external/torch-mlir/TorchPlugin/test/CMakeLists.txt vendored
View File

@ -1,21 +0,0 @@
configure_lit_site_cfg(
${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
MAIN_CONFIG
${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
)
set(TEST_DEPENDS
FileCheck count not
torch-mlir-opt
TorchMLIRPluginPythonModules
)
add_lit_testsuite(check-torch-mlir-plugin
"Running the torch-mlir-plugin regression tests"
${CMAKE_CURRENT_BINARY_DIR}
DEPENDS ${TEST_DEPENDS}
)
set_target_properties(check-torch-mlir-plugin PROPERTIES FOLDER "Tests")
add_lit_testsuites(TORCH_MLIR_PLUGIN ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${TEST_DEPENDS})

16
external/torch-mlir/TorchPlugin/test/extension_coexistence.py vendored
View File

@ -1,16 +0,0 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# Some checks that we can import the various extensions and libraries and
# not have symbol collisions or other goings on.
# RUN: %PYTHON %s
import sys
print(f"PYTHONPATH={sys.path}")
import mlir.ir
import torch_mlir
print("Extensions all loaded")

73
external/torch-mlir/TorchPlugin/test/lit.cfg.py vendored
View File

@ -1,73 +0,0 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See 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 = 'TORCH_MLIR_PLUGIN'
config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
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.torch_mlir_plugin_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.torch_mlir_plugin_obj_root, 'test')
config.substitutions.append(('%PATH%', config.environment['PATH']))
config.substitutions.append(('%shlibext', config.llvm_shlib_ext))
config.substitutions.append(('%PYTHON', config.python_executable))
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.torch_mlir_plugin_obj_root, 'test')
config.torch_mlir_tools_dir = os.path.join(config.torch_mlir_plugin_obj_root, 'bin')
# Tweak the PATH to include the tools dir.
llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
llvm_config.with_environment('PYTHONPATH', [
os.path.join(config.torch_mlir_python_packages_dir, 'torch_mlir'),
],
append_path=True)
tool_dirs = [config.torch_mlir_tools_dir, config.llvm_tools_dir]
tools = [
'torch-mlir-opt',
]
llvm_config.add_tool_substitutions(tools, tool_dirs)

56
external/torch-mlir/TorchPlugin/test/lit.site.cfg.py.in vendored
View File

@ -1,56 +0,0 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
@LIT_SITE_CFG_IN_HEADER@
import sys
config.host_triple = "@LLVM_HOST_TRIPLE@"
config.target_triple = "@TARGET_TRIPLE@"
config.llvm_src_root = "@LLVM_SOURCE_DIR@"
config.llvm_obj_root = "@LLVM_BINARY_DIR@"
# TODO: Fix tools dir to find FileCheck.
#config.llvm_tools_dir = "@LLVM_TOOLS_DIR@"
config.llvm_tools_dir = "@LLVM_BINARY_DIR@/bin"
config.llvm_lib_dir = "@LLVM_LIBRARY_DIR@"
config.llvm_shlib_dir = "@SHLIBDIR@"
config.llvm_shlib_ext = "@SHLIBEXT@"
config.llvm_exe_ext = "@EXEEXT@"
config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
config.python_executable = "@Python3_EXECUTABLE@"
config.gold_executable = "@GOLD_EXECUTABLE@"
config.ld64_executable = "@LD64_EXECUTABLE@"
config.enable_shared = @ENABLE_SHARED@
config.enable_assertions = @ENABLE_ASSERTIONS@
config.targets_to_build = "@TARGETS_TO_BUILD@"
config.native_target = "@LLVM_NATIVE_ARCH@"
config.llvm_bindings = "@LLVM_BINDINGS@".split(' ')
config.host_os = "@HOST_OS@"
config.host_cc = "@HOST_CC@"
config.host_cxx = "@HOST_CXX@"
# Note: ldflags can contain double-quoted paths, so must use single quotes here.
config.host_ldflags = '@HOST_LDFLAGS@'
config.llvm_use_sanitizer = "@LLVM_USE_SANITIZER@"
config.llvm_host_triple = '@LLVM_HOST_TRIPLE@'
config.host_arch = "@HOST_ARCH@"
config.torch_mlir_plugin_src_root = "@CMAKE_SOURCE_DIR@"
config.torch_mlir_plugin_obj_root = "@CMAKE_BINARY_DIR@"
config.torch_mlir_python_packages_dir = "@TORCH_MLIR_PYTHON_PACKAGES_DIR@"
# Support substitution of the tools_dir with user parameters. This is
# used when we can't determine the tool dir at configuration time.
try:
config.llvm_tools_dir = config.llvm_tools_dir % lit_config.params
config.llvm_shlib_dir = config.llvm_shlib_dir % lit_config.params
except KeyError:
e = sys.exc_info()[1]
key, = e.args
lit_config.fatal("unable to find %r parameter, use '--param=%s=VALUE'" % (key,key))
import lit.llvm
lit.llvm.initialize(lit_config, config)
# Let the main config do the real work.
lit_config.load_config(config, "@CMAKE_CURRENT_SOURCE_DIR@/lit.cfg.py")

9
external/torch-mlir/build_tools/build_standalone.sh vendored
View File

@ -15,11 +15,14 @@ llvm_project_dir="$project_dir/../llvm-project"
build_dir="$project_dir/build"
cmake -GNinja -B"$build_dir" "$llvm_project_dir/llvm" \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_BUILD_TYPE=RelWithDebInfo \
-DCMAKE_C_FLAGS_RELWITHDEBINFO="-O2 -DNDEBUG -gline-tables-only" \
-DCMAKE_CXX_FLAGS_RELWITHDEBINFO="-O2 -DNDEBUG -gline-tables-only" \
-DLLVM_ENABLE_PROJECTS=mlir \
-DLLVM_EXTERNAL_PROJECTS=torch-mlir \
-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR="$project_dir" \
-DMLIR_ENABLE_BINDINGS_PYTHON=ON
-DMLIR_ENABLE_BINDINGS_PYTHON=ON \
-DLLVM_ENABLE_ASSERTIONS=ON
cd "$build_dir"
ninja tools/torch-mlir/all check-torch-mlir check-torch-mlir-plugin
ninja tools/torch-mlir/all check-torch-mlir

3
external/torch-mlir/build_tools/update_torch_ods.sh vendored
View File

@ -8,6 +8,7 @@ torch_ir_dir="${src_dir}/include/torch-mlir/Dialect/Torch/IR"
python_packages_dir="${build_dir}/tools/torch-mlir/python_packages"
#ninja -C "${build_dir}"
PYTHONPATH="${python_packages_dir}/torch_mlir" python -m torch_mlir_utils.codegen.torch_ods_gen \
PYTHONPATH="${python_packages_dir}/torch_mlir" python \
-m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen \
--torch_ir_dir="${torch_ir_dir}" \
--debug_registry_dump="${torch_ir_dir}/JITOperatorRegistryDump.txt"

2
external/torch-mlir/include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td vendored
View File

@ -11,7 +11,7 @@
//
// This file is automatically generated. Please do not edit.
// Generated via:
// python -m torch_mlir_utils.codegen.torch_ods_gen
// python -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen
//
//===----------------------------------------------------------------------===//

2
external/torch-mlir/include/torch-mlir/Dialect/Torch/IR/GeneratedPrimOps.td vendored
View File

@ -11,7 +11,7 @@
//
// This file is automatically generated. Please do not edit.
// Generated via:
// python -m torch_mlir_utils.codegen.torch_ods_gen
// python -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen
//
//===----------------------------------------------------------------------===//

2
external/torch-mlir/include/torch-mlir/Dialect/Torch/IR/GeneratedQuantizedOps.td vendored
View File

@ -11,7 +11,7 @@
//
// This file is automatically generated. Please do not edit.
// Generated via:
// python -m torch_mlir_utils.codegen.torch_ods_gen
// python -m torch_mlir.dialects.torch.importer.jit_ir.build_tools.torch_ods_gen
//
//===----------------------------------------------------------------------===//

36
external/torch-mlir/python/CMakeLists.txt vendored
View File

@ -1,5 +1,13 @@
include(AddMLIRPython)
# The directory at which the Python import tree begins.
# See documentation for `declare_mlir_python_sources`'s ROOT_DIR
# argument.
set(TORCH_MLIR_PYTHON_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/torch_mlir")
# We vendor our own MLIR instance in the `torch_mlir` namespace.
add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=torch_mlir.")
################################################################################
# Sources
################################################################################
@ -8,14 +16,15 @@ declare_mlir_python_sources(TorchMLIRPythonSources)
declare_mlir_python_sources(TorchMLIRPythonExtensions)
declare_mlir_python_sources(TorchMLIRPythonSources.Dialects
ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
ADD_TO_PARENT TorchMLIRPythonSources
)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT TorchMLIRPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
TD_FILE dialects/TorchBinding.td
SOURCES dialects/torch.py
SOURCES dialects/torch/__init__.py
DIALECT_NAME torch
)
@ -34,6 +43,14 @@ declare_mlir_python_extension(TorchMLIRPythonExtensions.Main
LLVMSupport
)
################################################################################
# Optionally handle JIT IR importer.
################################################################################
if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
add_subdirectory(torch_mlir/dialects/torch/importer/jit_ir)
endif()
################################################################################
# Generate packages and shared library
# Downstreams typically will not use these, but they are useful for local
@ -52,16 +69,23 @@ set(_source_components
add_mlir_python_common_capi_library(TorchMLIRAggregateCAPI
INSTALL_COMPONENT TorchMLIRPythonModules
INSTALL_DESTINATION python_packages/torch_mlir/mlir/_mlir_libs
OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/mlir/_mlir_libs"
INSTALL_DESTINATION python_packages/torch_mlir/torch_mlir/_mlir_libs
OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs"
RELATIVE_INSTALL_ROOT "../../../.."
DECLARED_SOURCES ${_source_components}
)
add_mlir_python_modules(TorchMLIRPythonModules
ROOT_PREFIX "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/mlir"
INSTALL_PREFIX "python_packages/torch_mlir/mlir"
ROOT_PREFIX "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir"
INSTALL_PREFIX "python_packages/torch_mlir/torch_mlir"
DECLARED_SOURCES ${_source_components}
COMMON_CAPI_LINK_LIBS
TorchMLIRAggregateCAPI
)
# TODO: Find a cleaner way to do this.
# Can we build the JIT IR importer with `declare_mlir_python_extension`?
# Then it would "just work".
if(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER)
add_dependencies(TorchMLIRPythonModules TorchMLIRJITIRImporter)
endif()

2
external/torch-mlir/python/TorchMLIRModule.cpp vendored
View File

@ -17,7 +17,7 @@ PYBIND11_MODULE(_torchMlir, m) {
m.doc() = "torch-mlir main python extension";
m.def(
"register_torch_dialect",
"register_dialect",
[](MlirContext context, bool load) {
MlirDialectHandle handle = mlirGetDialectHandle__torch__();
mlirDialectHandleRegisterDialect(handle, context);

0
external/torch-mlir/python/mlir/dialects/TorchBinding.td → external/torch-mlir/python/torch_mlir/dialects/TorchBinding.td vendored
View File

4
external/torch-mlir/python/mlir/dialects/torch.py → external/torch-mlir/python/torch_mlir/dialects/torch/__init__.py vendored
View File

@ -2,5 +2,5 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from ._torch_ops_gen import *
from .._mlir_libs._torchMlir import register_torch_dialect
from .._torch_ops_gen import *
from ..._mlir_libs._torchMlir import register_dialect

0
external/torch-mlir/TorchPlugin/python/torch_mlir_utils/codegen/__init__.py → external/torch-mlir/python/torch_mlir/dialects/torch/importer/__init__.py vendored
View File

25
external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/CMakeLists.txt vendored 100644
View File

@ -0,0 +1,25 @@
#-------------------------------------------------------------------------------
# Setup PyTorch
#-------------------------------------------------------------------------------
list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules")
include(TorchMLIRPyTorch)
TorchMLIRProbeForPyTorchInstall()
find_package(Torch 1.8 REQUIRED)
TorchMLIRConfigurePyTorch()
#-------------------------------------------------------------------------------
# Subdirectories
#-------------------------------------------------------------------------------
add_subdirectory(csrc)
## Declare the sources of the Python module.
declare_mlir_python_sources(TorchMLIRPythonSources.JitIRImporter
ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
ADD_TO_PARENT TorchMLIRPythonSources
SOURCES_GLOB
dialects/torch/importer/jit_ir/*.py
)

2
external/torch-mlir/TorchPlugin/python/torch_mlir/__init__.py → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/__init__.py vendored
View File

@ -10,7 +10,7 @@ import torch
# Our native extension is not self-contained. It references libraries which
# must come in via the above first.
from _torch_mlir import *
from ....._mlir_libs._jit_ir_importer import *
__all__ = [

0
external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/__init__.py vendored 100644
View File

2
external/torch-mlir/TorchPlugin/python/torch_mlir_utils/codegen/torch_ods_gen.py → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py vendored
View File

@ -20,7 +20,7 @@ import textwrap
import traceback
# Note that this utility exists only in the c-extension.
from _torch_mlir import get_registered_ops # pytype: disable=import-error
from torch_mlir._mlir_libs._jit_ir_importer import get_registered_ops # pytype: disable=import-error
class TextEmitter:

0
external/torch-mlir/TorchPlugin/cmake/modules/TorchMLIRPyTorch.cmake → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/cmake/modules/TorchMLIRPyTorch.cmake vendored
View File

19
external/torch-mlir/TorchPlugin/csrc/CMakeLists.txt → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/CMakeLists.txt vendored
View File

@ -1,5 +1,3 @@
# TODO: Add a way to configure MLIR_PYTHON_PACKAGE_PREFIX
# Sharp edge: Torch extensions need to use the same pybind11 that torch
# was compiled with, or else there will be issues in cross module exception
# handling (which will abort instead of raise). We circumvent the possibility
@ -12,18 +10,18 @@ include_directories(BEFORE
)
link_directories("${TORCH_INSTALL_PREFIX}/lib")
add_library(TorchMLIRTorchPlugin SHARED
add_library(TorchMLIRJITIRImporter SHARED
class_annotator.cpp
get_registered_ops.cpp
function_importer.cpp
module_builder.cpp
node_importer.cpp
ivalue_importer.cpp
python_bindings.cpp
init_python_bindings.cpp
torch_to_mlir_utils.cpp
)
target_link_libraries(TorchMLIRTorchPlugin
target_link_libraries(TorchMLIRJITIRImporter
TorchMLIRAggregateCAPI
${TORCH_LIBRARIES}
${Python3_LIBRARIES}
@ -31,14 +29,15 @@ target_link_libraries(TorchMLIRTorchPlugin
)
message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS}")
set_target_properties(TorchMLIRTorchPlugin PROPERTIES
LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir"
OUTPUT_NAME _torch_mlir
set_target_properties(TorchMLIRJITIRImporter PROPERTIES
LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs"
OUTPUT_NAME _jit_ir_importer
PREFIX "${PYTHON_MODULE_PREFIX}"
SUFFIX "${PYTHON_MODULE_EXTENSION}"
CXX_VISIBILITY_PRESET "hidden"
COMPILE_FLAGS "${TORCH_CXXFLAGS}"
)
mlir_python_setup_extension_rpath(TorchMLIRJITIRImporter)
torch_mlir_python_target_compile_options(TorchMLIRTorchPlugin)
mlir_check_all_link_libraries(TorchMLIRTorchPlugin)
torch_mlir_python_target_compile_options(TorchMLIRJITIRImporter)
mlir_check_all_link_libraries(TorchMLIRJITIRImporter)

2
external/torch-mlir/TorchPlugin/csrc/class_annotator.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/class_annotator.cpp vendored
View File

@ -1,7 +1,7 @@
//===- class_annotator.cpp ------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/class_annotator.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/class_annotator.h vendored
View File

@ -1,7 +1,7 @@
//===- class_annotator.h ----------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
// Utilities for annotating Torch `c10::ClassType`
@ -18,8 +18,8 @@
// of implementation.
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_CLASS_ANNOTATOR_H
#define TORCHMLIRPLUGIN_CSRC_CLASS_ANNOTATOR_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_CLASS_ANNOTATOR_H
#define TORCHMLIRJITIRIMPORTER_CSRC_CLASS_ANNOTATOR_H
#include "pybind.h"
@ -192,4 +192,4 @@ void initClassAnnotatorBindings(py::module &m);
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_CLASS_ANNOTATOR_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_CLASS_ANNOTATOR_H

2
external/torch-mlir/TorchPlugin/csrc/function_importer.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/function_importer.cpp vendored
View File

@ -1,7 +1,7 @@
//===- function_importer.cpp ----------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/function_importer.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/function_importer.h vendored
View File

@ -1,12 +1,12 @@
//===- function_importer.h --------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_FUNCTION_IMPORTER_H
#define TORCHMLIRPLUGIN_CSRC_FUNCTION_IMPORTER_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_FUNCTION_IMPORTER_H
#define TORCHMLIRJITIRIMPORTER_CSRC_FUNCTION_IMPORTER_H
#include <memory>
@ -44,4 +44,4 @@ MlirOperation importJitFunctionAsFuncOp(
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_FUNCTION_IMPORTER_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_FUNCTION_IMPORTER_H

2
external/torch-mlir/TorchPlugin/csrc/get_registered_ops.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/get_registered_ops.cpp vendored
View File

@ -1,7 +1,7 @@
//===- get_registered_ops.cpp ---------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/get_registered_ops.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/get_registered_ops.h vendored
View File

@ -1,7 +1,7 @@
//===- get_registered_ops.h -------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
//
@ -10,8 +10,8 @@
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_GETREGISTEREDOPS_H
#define TORCHMLIRPLUGIN_CSRC_GETREGISTEREDOPS_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_GETREGISTEREDOPS_H
#define TORCHMLIRJITIRIMPORTER_CSRC_GETREGISTEREDOPS_H
#include "pybind.h"
@ -21,4 +21,4 @@ void initGetRegisteredOpsBindings(py::module &m);
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_GETREGISTEREDOPS_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_GETREGISTEREDOPS_H

6
external/torch-mlir/TorchPlugin/csrc/python_bindings.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/init_python_bindings.cpp vendored
View File

@ -1,11 +1,11 @@
//===- python_bindings.cpp --------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
// This is the top-level entry point for the MLIR <-> PyTorch bridge.
// This is the top-level entry point for the JIT IR -> MLIR importer.
#include <ATen/core/dispatch/Dispatcher.h>
@ -15,7 +15,7 @@
using namespace torch_mlir;
PYBIND11_MODULE(_torch_mlir, m) {
PYBIND11_MODULE(_jit_ir_importer, m) {
ModuleBuilder::bind(m);
initClassAnnotatorBindings(m);
initGetRegisteredOpsBindings(m);

2
external/torch-mlir/TorchPlugin/csrc/ivalue_importer.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/ivalue_importer.cpp vendored
View File

@ -1,7 +1,7 @@
//===- ivalue_importer.cpp ------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/ivalue_importer.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/ivalue_importer.h vendored
View File

@ -1,12 +1,12 @@
//===- ivalue_importer.h ----------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_IVALUE_IMPORTER_H
#define TORCHMLIRPLUGIN_CSRC_IVALUE_IMPORTER_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_IVALUE_IMPORTER_H
#define TORCHMLIRJITIRIMPORTER_CSRC_IVALUE_IMPORTER_H
#include <memory>
@ -28,4 +28,4 @@ void importIValue(c10::IValue ivalue, MlirBlock block, MlirContext context,
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_IVALUE_IMPORTER_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_IVALUE_IMPORTER_H

8
external/torch-mlir/TorchPlugin/csrc/mlir_utils.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/mlir_utils.h vendored
View File

@ -1,12 +1,12 @@
//===- mlir_utils.h ---------------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_MLIR_UTILS_H
#define TORCHMLIRPLUGIN_CSRC_MLIR_UTILS_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_MLIR_UTILS_H
#define TORCHMLIRJITIRIMPORTER_CSRC_MLIR_UTILS_H
#include <cstring>
#include <string>
@ -111,4 +111,4 @@ MlirOperation createMlirOperationAtEnd(MlirBlock block, std::string name,
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_MLIR_UTILS_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_MLIR_UTILS_H

2
external/torch-mlir/TorchPlugin/csrc/module_builder.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/module_builder.cpp vendored
View File

@ -1,7 +1,7 @@
//===- module_builder.cpp -------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/module_builder.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/module_builder.h vendored
View File

@ -1,12 +1,12 @@
//===- module_builder.h -----------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_BUILDER_H
#define TORCHMLIRPLUGIN_CSRC_BUILDER_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_BUILDER_H
#define TORCHMLIRJITIRIMPORTER_CSRC_BUILDER_H
#include "pybind.h"
@ -62,4 +62,4 @@ private:
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_C10_DISPATCH_MODULE_BUILDER_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_C10_DISPATCH_MODULE_BUILDER_H

2
external/torch-mlir/TorchPlugin/csrc/node_importer.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.cpp vendored
View File

@ -1,7 +1,7 @@
//===- node_importer.cpp --------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/node_importer.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/node_importer.h vendored
View File

@ -1,12 +1,12 @@
//===- node_importer.h ------------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_NODE_IMPORTER_H
#define TORCHMLIRPLUGIN_CSRC_NODE_IMPORTER_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_NODE_IMPORTER_H
#define TORCHMLIRJITIRIMPORTER_CSRC_NODE_IMPORTER_H
#include <memory>
@ -27,4 +27,4 @@ MlirBlock importBlock(MlirContext context, torch::jit::Block *jitBlock,
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_NODE_IMPORTER_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_NODE_IMPORTER_H

8
external/torch-mlir/TorchPlugin/csrc/pybind.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/pybind.h vendored
View File

@ -1,7 +1,7 @@
//===- module_builder.h -----------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
// Includes Torch-specific pybind and associated helpers.
@ -9,8 +9,8 @@
// directly).
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_PYBIND_H
#define TORCHMLIRPLUGIN_CSRC_PYBIND_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_PYBIND_H
#define TORCHMLIRJITIRIMPORTER_CSRC_PYBIND_H
#include <torch/csrc/utils/pybind.h>
@ -25,4 +25,4 @@ public:
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_PYBIND_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_PYBIND_H

2
external/torch-mlir/TorchPlugin/csrc/torch_to_mlir_utils.cpp → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.cpp vendored
View File

@ -1,7 +1,7 @@
//===- torch_to_mlir_utils.cpp --------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//

8
external/torch-mlir/TorchPlugin/csrc/torch_to_mlir_utils.h → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.h vendored
View File

@ -1,12 +1,12 @@
//===- torch_to_mlir_utils.h ------------------------------------*- C++ -*-===//
//
// This file is licensed under a pytorch-style license
// See LICENSE for license information.
// See LICENSE.pytorch for license information.
//
//===----------------------------------------------------------------------===//
#ifndef TORCHMLIRPLUGIN_CSRC_TORCH_TO_MLIR_UTILS_H
#define TORCHMLIRPLUGIN_CSRC_TORCH_TO_MLIR_UTILS_H
#ifndef TORCHMLIRJITIRIMPORTER_CSRC_TORCH_TO_MLIR_UTILS_H
#define TORCHMLIRJITIRIMPORTER_CSRC_TORCH_TO_MLIR_UTILS_H
#include <memory>
@ -76,4 +76,4 @@ MlirOperation createOperationFromSchema(MlirBlock appendToBlock,
} // namespace torch_mlir
#endif // TORCHMLIRPLUGIN_CSRC_TORCH_TO_MLIR_UTILS_H
#endif // TORCHMLIRJITIRIMPORTER_CSRC_TORCH_TO_MLIR_UTILS_H

9
external/torch-mlir/TorchPlugin/python/torch_mlir/torchscript_annotations.py → external/torch-mlir/python/torch_mlir/dialects/torch/importer/jit_ir/torchscript_annotations.py vendored
View File

@ -7,11 +7,12 @@ from typing import List, Optional, Tuple
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator
# Decorators
# Currently, these decorators are very low-level and map 1:1 with
# methods on `torch_mlir.ClassAnnotator`. Eventually, we expect there to
# methods on `ClassAnnotator`. Eventually, we expect there to
# be a more elaborate Python layer which allows all the different annotations
# to be expressed conveniently and gives clearer error reports when
# the annotations aren't acceptable.
@ -20,11 +21,11 @@ import torch_mlir
# we can use that module from code without C++ dependencies, which prevent us
# from interfacing the test framework across environments.
# Utilities for extracting decorated information into torch_mlir.ClassAnnotator.
# Utilities for extracting decorated information into ClassAnnotator.
def _recursively_extract_annotations(
module: torch.nn.Module, scripted: torch.jit.ScriptModule,
class_annotator: torch_mlir.ClassAnnotator):
class_annotator: ClassAnnotator):
assert module.__class__.__name__ == scripted.original_name or (
isinstance(module, torch.jit.RecursiveScriptModule) and module is
scripted), "script module does not come from specified module"
@ -49,7 +50,7 @@ def _recursively_extract_annotations(
def extract_annotations(program: torch.nn.Module,
scripted: torch.jit.ScriptModule,
class_annotator: torch_mlir.ClassAnnotator):
class_annotator: ClassAnnotator):
"""Populate the ClassAnnotator with annotations extracted from `program`."""
class_annotator.exportNone(scripted._c._type())
_recursively_extract_annotations(program, scripted, class_annotator)

1
external/torch-mlir/test/CMakeLists.txt vendored
View File

@ -1,5 +1,6 @@
llvm_canonicalize_cmake_booleans(
MLIR_ENABLE_BINDINGS_PYTHON
TORCH_MLIR_ENABLE_JIT_IR_IMPORTER
)
configure_lit_site_cfg(

1
external/torch-mlir/test/lit.site.cfg.py.in vendored
View File

@ -14,6 +14,7 @@ config.llvm_shlib_ext = "@SHLIBEXT@"
config.llvm_exe_ext = "@EXEEXT@"
config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
config.python_executable = sys.executable
config.enable_jit_ir_importer = @TORCH_MLIR_ENABLE_JIT_IR_IMPORTER@
import lit.llvm
lit.llvm.initialize(lit_config, config)

6
external/torch-mlir/TorchPlugin/test/get_registered_ops.py → external/torch-mlir/test/python/importer/jit_ir/get_registered_ops.py vendored
View File

@ -1,11 +1,11 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
# RUN: %PYTHON %s | FileCheck %s
import _torch_mlir
from torch_mlir._mlir_libs._jit_ir_importer import get_registered_ops
# This check is just for a built-in op that is unlikely to change (and is
# otherwise insignificant).
# CHECK: {'name': ('aten::mul', 'Tensor'), 'is_c10_op': True, 'is_vararg': False, 'is_varret': False, 'is_mutable': False, 'arguments': [{'name': 'self', 'type': 'Tensor', 'pytype': 'Tensor'}, {'name': 'other', 'type': 'Tensor', 'pytype': 'Tensor'}], 'returns': [{'name': '', 'type': 'Tensor', 'pytype': 'Tensor'}]}
print('\n\n'.join([repr(r) for r in _torch_mlir.get_registered_ops()]))
print('\n\n'.join([repr(r) for r in get_registered_ops()]))

0
external/torch-mlir/TorchPlugin/test/ivalue_import/README.md → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/README.md vendored
View File

9
external/torch-mlir/TorchPlugin/test/ivalue_import/annotations/arg-error.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/annotations/arg-error.py vendored
View File

@ -1,15 +1,14 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
@ -23,7 +22,7 @@ class TestModule(torch.nn.Module):
test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()
try:
annotator.annotateArgs(class_type, [], [])

9
external/torch-mlir/TorchPlugin/test/ivalue_import/annotations/arg-tensor-type-bound.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/annotations/arg-tensor-type-bound.py vendored
View File

@ -1,15 +1,14 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):
@ -20,7 +19,7 @@ class TestModule(torch.nn.Module):
test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()
# CHECK: func private @__torch__.TestModule.forward(
# CHECK-SAME: %arg0: !torch.nn.Module<"__torch__.TestModule">,

9
external/torch-mlir/TorchPlugin/test/ivalue_import/annotations/class-annotator-repr.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/annotations/class-annotator-repr.py vendored
View File

@ -1,15 +1,14 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class Submodule(torch.nn.Module):
@ -37,7 +36,7 @@ class TestModule(torch.nn.Module):
test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()
annotator.exportNone(class_type)

9
external/torch-mlir/TorchPlugin/test/ivalue_import/annotations/export-error.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/annotations/export-error.py vendored
View File

@ -1,15 +1,14 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):
@ -20,7 +19,7 @@ class TestModule(torch.nn.Module):
test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()
try:

9
external/torch-mlir/TorchPlugin/test/ivalue_import/annotations/export-recursive.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/annotations/export-recursive.py vendored
View File

@ -1,15 +1,14 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class Submodule(torch.nn.Module):
@ -32,7 +31,7 @@ class TestModule(torch.nn.Module):
test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()
# CHECK-LABEL: torch.class_type @__torch__.Submodule {
# CHECK: torch.attr "exported" : !torch.int

9
external/torch-mlir/TorchPlugin/test/ivalue_import/annotations/export.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/annotations/export.py vendored
View File

@ -1,15 +1,14 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):
@ -24,7 +23,7 @@ class TestModule(torch.nn.Module):
test_module = TestModule()
recursivescriptmodule = torch.jit.script(test_module)
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
class_type = recursivescriptmodule._c._type()
# CHECK-LABEL: torch.class_type @__torch__.TestModule {
# CHECK: torch.attr "exported" : !torch.int

6
external/torch-mlir/TorchPlugin/test/ivalue_import/dict.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/dict.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
from typing import Dict, Optional
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/functions-that-call-methods.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/functions-that-call-methods.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# Interesting test case, where a function calls a method.

6
external/torch-mlir/TorchPlugin/test/ivalue_import/functions.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/functions.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func private @__torch__.TestModule.forward
# CHECK-SAME: (%[[ARG0:.*]]: !torch.nn.Module<"__torch__.TestModule">, %[[ARG1:.*]]: !torch.tensor) -> !torch.tensor {

6
external/torch-mlir/TorchPlugin/test/ivalue_import/list.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/list.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/methods-derefine.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/methods-derefine.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()

6
external/torch-mlir/TorchPlugin/test/ivalue_import/methods-locations.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/methods-locations.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/methods.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/methods.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# Function names in the Torch compilation unit are systematic -- they

6
external/torch-mlir/TorchPlugin/test/ivalue_import/object-identity-error-submodule.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/object-identity-error-submodule.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: not %PYTHON %s 2>&1 | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class Submodule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/object-identity-error.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/object-identity-error.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: not %PYTHON %s 2>&1 | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/object-identity-torch-bug.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/object-identity-torch-bug.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# TorchScript doesn't model object identity correctly!!!
#

6
external/torch-mlir/TorchPlugin/test/ivalue_import/object-identity.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/object-identity.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/prim.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/prim.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/primitives.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/primitives.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/quantization.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/quantization.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/strings.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/strings.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/submodules-select.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/submodules-select.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class Submodule(torch.nn.Module):
def __init__(self, n):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/submodules.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/submodules.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class Submodule(torch.nn.Module):
def __init__(self, n):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/tensors.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/tensors.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

6
external/torch-mlir/TorchPlugin/test/ivalue_import/tuple.py → external/torch-mlir/test/python/importer/jit_ir/ivalue_import/tuple.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
class TestModule(torch.nn.Module):
def __init__(self):

2
external/torch-mlir/test/python/importer/jit_ir/lit.local.cfg vendored 100644
View File

@ -0,0 +1,2 @@
if not config.enable_jit_ir_importer:
config.unsupported = True

0
external/torch-mlir/TorchPlugin/test/node_import/README.md → external/torch-mlir/test/python/importer/jit_ir/node_import/README.md vendored
View File

6
external/torch-mlir/TorchPlugin/test/node_import/debug-info.py → external/torch-mlir/test/python/importer/jit_ir/node_import/debug-info.py vendored
View File

@ -1,13 +1,13 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func @__torch__.add3
# Note that line-level debug information for parts unannotated in the Torch

6
external/torch-mlir/TorchPlugin/test/node_import/dict.py → external/torch-mlir/test/python/importer/jit_ir/node_import/dict.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
import collections
from typing import Tuple, Optional, List, NamedTuple, Dict
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func @__torch__.dict_literal_empty() -> !torch.dict<!torch.str, !torch.tensor> {

6
external/torch-mlir/TorchPlugin/test/node_import/elif.py → external/torch-mlir/test/python/importer/jit_ir/node_import/elif.py vendored
View File

@ -1,13 +1,13 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: @__torch__.f
@mb.import_function

6
external/torch-mlir/TorchPlugin/test/node_import/errors.py → external/torch-mlir/test/python/importer/jit_ir/node_import/errors.py vendored
View File

@ -1,11 +1,11 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import enum
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
class Color(enum.Enum):
@ -15,7 +15,7 @@ class Color(enum.Enum):
# RUN: %PYTHON %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# To test errors, use a type that we don't support yet.
try:

6
external/torch-mlir/TorchPlugin/test/node_import/function-derefine.py → external/torch-mlir/test/python/importer/jit_ir/node_import/function-derefine.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
import typing
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func @__torch__.optional_return(
# CHECK-SAME: %[[ARG:.*]]: !torch.int) -> !torch.optional<!torch.int> {

6
external/torch-mlir/TorchPlugin/test/node_import/if.py → external/torch-mlir/test/python/importer/jit_ir/node_import/if.py vendored
View File

@ -1,13 +1,13 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# Note: The "if without else" case is handled by yielding None from the
# else branch and making all defined values optional, so no special handling

6
external/torch-mlir/TorchPlugin/test/node_import/list.py → external/torch-mlir/test/python/importer/jit_ir/node_import/list.py vendored
View File

@ -1,13 +1,13 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func @__torch__.f(
# CHECK-SAME: %[[T0:.*]]: !torch.tensor,

6
external/torch-mlir/TorchPlugin/test/node_import/loop.py → external/torch-mlir/test/python/importer/jit_ir/node_import/loop.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
import typing
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func @__torch__.prim_Loop_forlike(
# CHECK-SAME: %[[MAX_ITERATIONS:.*]]: !torch.int) -> !torch.float {

6
external/torch-mlir/TorchPlugin/test/node_import/prim.py → external/torch-mlir/test/python/importer/jit_ir/node_import/prim.py vendored
View File

@ -1,17 +1,17 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import typing
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
import typing
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK-LABEL: func @__torch__.prim_NumToTensor(

6
external/torch-mlir/TorchPlugin/test/node_import/tuple.py → external/torch-mlir/test/python/importer/jit_ir/node_import/tuple.py vendored
View File

@ -1,15 +1,15 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
import collections
from typing import Tuple, Optional, NamedTuple
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
NT = NamedTuple('NT', [('f1', Optional[torch.Tensor]),
('f2', Optional[torch.Tensor])])

6
external/torch-mlir/TorchPlugin/test/node_import/types-bool.py → external/torch-mlir/test/python/importer/jit_ir/node_import/types-bool.py vendored
View File

@ -1,13 +1,13 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK: @__torch__.returns_bool
@mb.import_function

6
external/torch-mlir/TorchPlugin/test/node_import/types-none.py → external/torch-mlir/test/python/importer/jit_ir/node_import/types-none.py vendored
View File

@ -1,13 +1,13 @@
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See LICENSE for license information.
# See LICENSE.pytorch for license information.
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ModuleBuilder
# RUN: %PYTHON %s | torch-mlir-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
# CHECK: @__torch__.returns_none
@mb.import_function

8
external/torch-mlir/test/python/smoketest.py vendored
View File

@ -1,7 +1,7 @@
# RUN: %PYTHON %s
import mlir.ir
from mlir.dialects import torch
import torch_mlir.ir
from torch_mlir.dialects import torch
with mlir.ir.Context() as ctx:
torch.register_torch_dialect(ctx)
with torch_mlir.ir.Context() as ctx:
torch.register_dialect(ctx)

1
python/CMakeLists.txt
View File

@ -111,7 +111,6 @@ add_mlir_python_modules(NPCOMPPythonModules
)
add_dependencies(NPCOMPPythonModules
NPCOMPPythonResources
TorchMLIRPluginPythonModules
)
################################################################################

6
python/test/annotations-sugar.py
View File

@ -6,9 +6,9 @@
import torch
import torch_mlir
from npcomp_torchscript.annotations import annotate_args, export
from torch_mlir.torchscript_annotations import extract_annotations
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator
from torch_mlir.dialects.torch.importer.jit_ir.torchscript_annotations import extract_annotations
class MmModule(torch.nn.Module):
def __init__(self):
@ -23,7 +23,7 @@ class MmModule(torch.nn.Module):
return torch.mm(lhs, rhs)
module = MmModule()
annotator = torch_mlir.ClassAnnotator()
annotator = ClassAnnotator()
extract_annotations(module, torch.jit.script(module), annotator)
print(annotator)

2
python/test/lit.cfg.py
View File

@ -18,7 +18,7 @@ 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.name = 'NPCOMP_PYTHON'
config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
if 'TEST_SRC_PATH' in os.environ:

8
python/torch_support/npcomp_torchscript_e2e_test_configs/npcomp_backend.py
View File

@ -11,13 +11,13 @@ import tempfile
import numpy as np
import torch
import torch_mlir
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
from torch_mlir.dialects.torch.importer.jit_ir.torchscript_annotations import extract_annotations
import npcomp
from npcomp.passmanager import PassManager
from npcomp.compiler.pytorch.backend import refjit
from npcomp.compiler.pytorch.backend.abc import NpcompBackend
from npcomp_torchscript.e2e_test.framework import TestConfig, Trace, TraceItem
from torch_mlir.torchscript_annotations import extract_annotations
def _recursively_convert_to_numpy(o: Any):
if isinstance(o, torch.Tensor):
@ -66,9 +66,9 @@ class NpcompBackendTestConfig(TestConfig):
self.backend = backend
def compile(self, program: torch.nn.Module) -> Any:
mb = torch_mlir.ModuleBuilder()
mb = ModuleBuilder()
scripted = torch.jit.script(program)
class_annotator = torch_mlir.ClassAnnotator()
class_annotator = ClassAnnotator()
extract_annotations(program, scripted, class_annotator)