This is an automated email from the ASF dual-hosted git repository.
echuraev pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e11913be06 [Target] Add target to all TVM callbacks (#14939)
e11913be06 is described below
commit e11913be06b30d806e6c09a6610b4e8fbaaa30a0
Author: Junru Shao <[email protected]>
AuthorDate: Wed May 24 23:10:41 2023 -0700
[Target] Add target to all TVM callbacks (#14939)
* [Target] Add target to all TVM callbacks
This PR adds an extra parameter `target` to all `tvm_callback_*` so that
the callback can decide its own behavior by querying which target to
compile against.
* fix lint
* fix lint
---
apps/ios_rpc/tests/ios_rpc_mobilenet.py | 27 +++++++++++-----------
apps/ios_rpc/tests/ios_rpc_test.py | 13 ++++++-----
apps/topi_recipe/broadcast/test_broadcast_map.py | 12 ++++------
apps/topi_recipe/conv/depthwise_conv2d_test.py | 13 +++++------
apps/topi_recipe/conv/test_conv2d_hwcn_map.py | 9 ++++----
apps/topi_recipe/reduce/test_reduce_map.py | 10 ++++----
apps/topi_recipe/rnn/lstm.py | 9 ++++----
apps/topi_recipe/rnn/matexp.py | 13 ++++++-----
jvm/core/src/test/scripts/test_add_gpu.py | 4 ++--
python/tvm/contrib/nvcc.py | 6 ++---
python/tvm/contrib/sdaccel.py | 10 ++++----
src/target/opt/build_cuda_on.cc | 4 ++--
src/target/source/codegen_aocl.cc | 2 +-
src/target/source/codegen_metal.cc | 2 +-
src/target/source/codegen_opencl.cc | 2 +-
src/target/source/codegen_vhls.cc | 5 ++--
src/target/spirv/spirv_utils.cc | 2 +-
tests/python/integration/test_ewise.py | 5 ++--
tests/python/integration/test_ewise_fpga.py | 3 +--
.../test_tir_transform_inject_ptx_async_copy.py | 2 +-
20 files changed, 76 insertions(+), 77 deletions(-)
diff --git a/apps/ios_rpc/tests/ios_rpc_mobilenet.py
b/apps/ios_rpc/tests/ios_rpc_mobilenet.py
index b90b459280..0c958a07d8 100644
--- a/apps/ios_rpc/tests/ios_rpc_mobilenet.py
+++ b/apps/ios_rpc/tests/ios_rpc_mobilenet.py
@@ -15,24 +15,24 @@
# specific language governing permissions and limitations
# under the License.
-import tvm
-from tvm import rpc, relay
-from tvm.contrib.download import download_testdata
-from tvm.relay.expr_functor import ExprMutator
-from tvm.relay import transform
-from tvm.relay.op.annotation import compiler_begin, compiler_end
-from tvm.relay.quantize.quantize import prerequisite_optimize
-from tvm.contrib import utils, xcode, graph_executor, coreml_runtime
-from tvm.contrib.target import coreml as _coreml
-
+import argparse
import os
import re
import sys
+
+import coremltools
import numpy as np
+import tvm
from mxnet import gluon
from PIL import Image
-import coremltools
-import argparse
+from tvm import relay, rpc
+from tvm.contrib import coreml_runtime, graph_executor, utils, xcode
+from tvm.contrib.download import download_testdata
+from tvm.contrib.target import coreml as _coreml
+from tvm.relay import transform
+from tvm.relay.expr_functor import ExprMutator
+from tvm.relay.op.annotation import compiler_begin, compiler_end
+from tvm.relay.quantize.quantize import prerequisite_optimize
# Change target configuration, this is setting for iphone6s
# arch = "x86_64"
@@ -43,9 +43,10 @@ target_host = "llvm -mtriple=%s-apple-darwin" % arch
MODES = {"proxy": rpc.connect, "tracker": rpc.connect_tracker, "standalone":
rpc.connect}
+
# override metal compiler to compile to iphone
@tvm.register_func("tvm_callback_metal_compile")
-def compile_metal(src):
+def compile_metal(src, target):
return xcode.compile_metal(src, sdk=sdk)
diff --git a/apps/ios_rpc/tests/ios_rpc_test.py
b/apps/ios_rpc/tests/ios_rpc_test.py
index 94340dcd4e..78f1d3a6c5 100644
--- a/apps/ios_rpc/tests/ios_rpc_test.py
+++ b/apps/ios_rpc/tests/ios_rpc_test.py
@@ -20,15 +20,15 @@ To use it, start a rpc proxy with "python -m
tvm.exec.rpc_proxy".
And configure the proxy host field as commented.
"""
-import tvm
-from tvm import te
+import argparse
import os
import re
import sys
-from tvm import rpc
-from tvm.contrib import utils, xcode
+
import numpy as np
-import argparse
+import tvm
+from tvm import rpc, te
+from tvm.contrib import utils, xcode
# Change target configuration, this is setting for iphone6s
arch = "arm64"
@@ -37,9 +37,10 @@ target = "llvm -mtriple=%s-apple-darwin" % arch
MODES = {"proxy": rpc.connect, "tracker": rpc.connect_tracker, "standalone":
rpc.connect}
+
# override metal compiler to compile to iphone
@tvm.register_func("tvm_callback_metal_compile")
-def compile_metal(src):
+def compile_metal(src, target):
return xcode.compile_metal(src, sdk=sdk)
diff --git a/apps/topi_recipe/broadcast/test_broadcast_map.py
b/apps/topi_recipe/broadcast/test_broadcast_map.py
index 4840a292d4..4d1416c41d 100644
--- a/apps/topi_recipe/broadcast/test_broadcast_map.py
+++ b/apps/topi_recipe/broadcast/test_broadcast_map.py
@@ -15,20 +15,18 @@
# specific language governing permissions and limitations
# under the License.
import os
+
+import numpy as np
import tvm
-from tvm import te
+from tvm import te, topi
from tvm.contrib import nvcc
-import numpy as np
-
-from tvm import topi
-
TASK = "reduce_map"
USE_MANUAL_CODE = False
@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target):
ptx = nvcc.compile_cuda(code, target_format="ptx")
return ptx
@@ -39,7 +37,7 @@ def write_code(code, fname):
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, target):
if not os.path.exists("perf"):
os.mkdir("perf")
write_code(code, "perf/%s_generated.cu" % TASK)
diff --git a/apps/topi_recipe/conv/depthwise_conv2d_test.py
b/apps/topi_recipe/conv/depthwise_conv2d_test.py
index 5ec205df52..eba0b4d38d 100644
--- a/apps/topi_recipe/conv/depthwise_conv2d_test.py
+++ b/apps/topi_recipe/conv/depthwise_conv2d_test.py
@@ -15,25 +15,24 @@
# specific language governing permissions and limitations
# under the License.
import os
-import tvm
-from tvm import te
+
import numpy as np
+import tvm
from scipy import signal
+from tvm import te, topi
from tvm.contrib import nvcc
-
-from tvm import topi
-from tvm.topi.utils import get_const_tuple
from tvm.topi.cuda.depthwise_conv2d import (
schedule_depthwise_conv2d_nchw,
schedule_depthwise_conv2d_nhwc,
)
+from tvm.topi.utils import get_const_tuple
TASK = "depthwise_conv2d"
USE_MANUAL_CODE = False
@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target):
ptx = nvcc.compile_cuda(code, target_format="ptx")
return ptx
@@ -44,7 +43,7 @@ def write_code(code, fname):
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, target):
if not os.path.exists("perf"):
os.mkdir("perf")
write_code(code, "perf/%s_generated.cu" % TASK)
diff --git a/apps/topi_recipe/conv/test_conv2d_hwcn_map.py
b/apps/topi_recipe/conv/test_conv2d_hwcn_map.py
index 6b239edb65..5d85d68ac3 100644
--- a/apps/topi_recipe/conv/test_conv2d_hwcn_map.py
+++ b/apps/topi_recipe/conv/test_conv2d_hwcn_map.py
@@ -16,12 +16,11 @@
# under the License.
"""Example code to do convolution."""
import os
+
import numpy as np
-import scipy.signal
import tvm
-from tvm import te
+from tvm import te, topi
from tvm.contrib import nvcc
-from tvm import topi
from tvm.topi.utils import get_const_tuple
TASK = "conv2d_hwcn_map"
@@ -29,7 +28,7 @@ USE_MANUAL_CODE = False
@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target):
ptx = nvcc.compile_cuda(code, target_format="ptx")
return ptx
@@ -40,7 +39,7 @@ def write_code(code, fname):
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, target):
if not os.path.exists("perf"):
os.mkdir("perf")
write_code(code, "perf/%s_generated.cu" % TASK)
diff --git a/apps/topi_recipe/reduce/test_reduce_map.py
b/apps/topi_recipe/reduce/test_reduce_map.py
index f8d63e2d19..c3d3096986 100644
--- a/apps/topi_recipe/reduce/test_reduce_map.py
+++ b/apps/topi_recipe/reduce/test_reduce_map.py
@@ -15,13 +15,11 @@
# specific language governing permissions and limitations
# under the License.
import os
+
+import numpy as np
import tvm
-from tvm import te
+from tvm import te, topi
from tvm.contrib import nvcc
-import numpy as np
-
-from tvm import topi
-
TASK = "reduce_map"
USE_MANUAL_CODE = False
@@ -33,7 +31,7 @@ def write_code(code, fname):
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, target):
if not os.path.exists("perf"):
os.mkdir("perf")
write_code(code, "perf/%s_generated.cu" % TASK)
diff --git a/apps/topi_recipe/rnn/lstm.py b/apps/topi_recipe/rnn/lstm.py
index bb9e31c5b2..11d072c175 100644
--- a/apps/topi_recipe/rnn/lstm.py
+++ b/apps/topi_recipe/rnn/lstm.py
@@ -15,11 +15,12 @@
# specific language governing permissions and limitations
# under the License.
"""LSTM Example, still work in progress.."""
+import os
+
+import numpy as np
import tvm
from tvm import te
-import os
from tvm.contrib import nvcc
-import numpy as np
# Quick knobs
TASK = "lstm"
@@ -31,7 +32,7 @@ UNROLL_WLOAD = True
@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target):
"""Use nvcc compiler for better perf."""
ptx = nvcc.compile_cuda(code, target_format="ptx")
return ptx
@@ -43,7 +44,7 @@ def write_code(code, fname):
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, target):
if not os.path.exists("perf"):
os.mkdir("perf")
write_code(code, "perf/%s_generated.cu" % TASK)
diff --git a/apps/topi_recipe/rnn/matexp.py b/apps/topi_recipe/rnn/matexp.py
index 303f0ed80d..6a7d58cba1 100644
--- a/apps/topi_recipe/rnn/matexp.py
+++ b/apps/topi_recipe/rnn/matexp.py
@@ -23,13 +23,14 @@ which calculates the following recursion formula
X[t] = dot(X[t-1], W)
```
"""
+import argparse
+import os
+import time
+
+import numpy as np
import tvm
from tvm import te
-import time
-import os
-import argparse
from tvm.contrib import nvcc
-import numpy as np
# Quick knobs
TASK = "matexp"
@@ -40,7 +41,7 @@ SKIP_CHECK = False
@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target):
"""Use nvcc compiler for better perf."""
ptx = nvcc.compile_cuda(code, target_format="ptx")
return ptx
@@ -52,7 +53,7 @@ def write_code(code, fname):
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, target):
if not os.path.exists("perf"):
os.mkdir("perf")
write_code(code, "perf/%s_generated.cu" % TASK)
diff --git a/jvm/core/src/test/scripts/test_add_gpu.py
b/jvm/core/src/test/scripts/test_add_gpu.py
index 21fd9edc06..0eea5671ba 100644
--- a/jvm/core/src/test/scripts/test_add_gpu.py
+++ b/jvm/core/src/test/scripts/test_add_gpu.py
@@ -18,11 +18,11 @@ import os
import tvm
from tvm import te
-from tvm.contrib import cc, utils, nvcc
+from tvm.contrib import cc, nvcc, utils
@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target):
ptx = nvcc.compile_cuda(code, target_format="ptx")
return ptx
diff --git a/python/tvm/contrib/nvcc.py b/python/tvm/contrib/nvcc.py
index 8acd620252..643ad96c02 100644
--- a/python/tvm/contrib/nvcc.py
+++ b/python/tvm/contrib/nvcc.py
@@ -18,15 +18,15 @@
"""Utility to invoke nvcc compiler in the system"""
from __future__ import absolute_import as _abs
-import subprocess
import os
+import subprocess
import warnings
import tvm._ffi
from tvm.target import Target
-from . import utils
from .._ffi.base import py_str
+from . import utils
def compile_cuda(code, target_format="ptx", arch=None, options=None,
path_target=None):
@@ -184,7 +184,7 @@ def get_cuda_version(cuda_path=None):
@tvm._ffi.register_func
-def tvm_callback_cuda_compile(code):
+def tvm_callback_cuda_compile(code, target): # pylint: disable=unused-argument
"""use nvcc to generate fatbin code for better optimization"""
ptx = compile_cuda(code, target_format="fatbin")
return ptx
diff --git a/python/tvm/contrib/sdaccel.py b/python/tvm/contrib/sdaccel.py
index 930752c2bc..478436e3d5 100644
--- a/python/tvm/contrib/sdaccel.py
+++ b/python/tvm/contrib/sdaccel.py
@@ -15,15 +15,16 @@
# specific language governing permissions and limitations
# under the License.
"""Utility for Interacting with SDAccel Tools"""
-import subprocess
import os
+import subprocess
import tvm._ffi
+
from . import utils
@tvm._ffi.register_func("tvm_callback_sdaccel_compile")
-def compile_vhls(kernel_info, device_name):
+def compile_vhls(kernel_info, target):
"""Compile Vivado HLS code for SDAccel.
Parameters
@@ -32,14 +33,15 @@ def compile_vhls(kernel_info, device_name):
List of kernel information. The kernel information is a tuple of
function name and source code.
- device_name : str
- The name of the target device
+ target : tvm.target.Target
+ The compilation target
Return
------
xclbin : bytearray
The bytearray of the xclbin
"""
+ device_name = target.attrs.get("device", "")
tmp_dir = utils.tempdir()
sdk = os.environ.get("XILINX_SDX", None)
diff --git a/src/target/opt/build_cuda_on.cc b/src/target/opt/build_cuda_on.cc
index 3be7103314..1c0b5094ef 100644
--- a/src/target/opt/build_cuda_on.cc
+++ b/src/target/opt/build_cuda_on.cc
@@ -143,14 +143,14 @@ runtime::Module BuildCUDA(IRModule mod, Target target) {
std::string code = cg.Finish();
if (const auto* f = Registry::Get("tvm_callback_cuda_postproc")) {
- code = (*f)(code).operator std::string();
+ code = (*f)(code, target).operator std::string();
}
std::string fmt = "ptx";
std::string ptx;
const auto* f_enter = Registry::Get("target.TargetEnterScope");
(*f_enter)(target);
if (const auto* f = Registry::Get("tvm_callback_cuda_compile")) {
- ptx = (*f)(code).operator std::string();
+ ptx = (*f)(code, target).operator std::string();
// Dirty matching to check PTX vs cubin.
// TODO(tqchen) more reliable checks
if (ptx[0] != '/') fmt = "cubin";
diff --git a/src/target/source/codegen_aocl.cc
b/src/target/source/codegen_aocl.cc
index 17e38e9af6..700d85b4cc 100644
--- a/src/target/source/codegen_aocl.cc
+++ b/src/target/source/codegen_aocl.cc
@@ -51,7 +51,7 @@ runtime::Module BuildAOCL(IRModule mod, Target target, bool
emulation) {
std::string code = cg.Finish();
if (const auto* f = Registry::Get("tvm_callback_opencl_postproc")) {
- code = (*f)(code).operator std::string();
+ code = (*f)(code, target).operator std::string();
}
// Write a .cl file.
diff --git a/src/target/source/codegen_metal.cc
b/src/target/source/codegen_metal.cc
index 9288c94e3d..bd2b930166 100644
--- a/src/target/source/codegen_metal.cc
+++ b/src/target/source/codegen_metal.cc
@@ -365,7 +365,7 @@ runtime::Module BuildMetal(IRModule mod, Target target) {
std::string fsource = cg.Finish();
source_maker << fsource << "\n";
if (fmetal_compile) {
- fsource = (*fmetal_compile)(fsource).operator std::string();
+ fsource = (*fmetal_compile)(fsource, target).operator std::string();
}
smap[func_name] = fsource;
}
diff --git a/src/target/source/codegen_opencl.cc
b/src/target/source/codegen_opencl.cc
index 14fbc4d840..61a8ee8a57 100644
--- a/src/target/source/codegen_opencl.cc
+++ b/src/target/source/codegen_opencl.cc
@@ -607,7 +607,7 @@ runtime::Module BuildOpenCL(IRModule mod, Target target) {
cg.AddFunction(f);
std::string fsource = cg.Finish();
if (fpostproc) {
- fsource = (*fpostproc)(fsource).operator std::string();
+ fsource = (*fpostproc)(fsource, target).operator std::string();
}
code << fsource;
}
diff --git a/src/target/source/codegen_vhls.cc
b/src/target/source/codegen_vhls.cc
index 3ae3fb773d..8463d6ac41 100644
--- a/src/target/source/codegen_vhls.cc
+++ b/src/target/source/codegen_vhls.cc
@@ -167,7 +167,7 @@ runtime::Module BuildSDAccel(IRModule mod, Target target) {
cg.AddFunction(f);
std::string code = cg.Finish();
if (const auto* f = runtime::Registry::Get("tvm_callback_vhls_postproc")) {
- code = (*f)(code).operator std::string();
+ code = (*f)(code, target).operator std::string();
}
auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
@@ -178,8 +178,7 @@ runtime::Module BuildSDAccel(IRModule mod, Target target) {
std::string xclbin;
if (const auto* f = Registry::Get("tvm_callback_sdaccel_compile")) {
- String device = target->GetAttr<String>("device", "").value();
- xclbin = (*f)(kernel_info, device).operator std::string();
+ xclbin = (*f)(kernel_info, target).operator std::string();
} else {
LOG(FATAL) << "Cannot compile Vivado HLS code.";
}
diff --git a/src/target/spirv/spirv_utils.cc b/src/target/spirv/spirv_utils.cc
index 2a9110d871..1bfbf63a81 100644
--- a/src/target/spirv/spirv_utils.cc
+++ b/src/target/spirv/spirv_utils.cc
@@ -161,7 +161,7 @@ std::pair<std::unordered_map<std::string,
runtime::SPIRVShader>, std::string> Lo
TVMByteArray arr;
arr.data = reinterpret_cast<const char*>(dmlc::BeginPtr(shader.data));
arr.size = shader.data.size() * sizeof(uint32_t);
- std::string transformed = (*postproc)(arr);
+ std::string transformed = (*postproc)(arr, target);
ICHECK_EQ(transformed.length() % 4U, 0U);
shader.data.resize(transformed.size() / 4U);
std::copy(transformed.begin(), transformed.end(),
diff --git a/tests/python/integration/test_ewise.py
b/tests/python/integration/test_ewise.py
index 8bfa6b1717..d1d2b9902c 100644
--- a/tests/python/integration/test_ewise.py
+++ b/tests/python/integration/test_ewise.py
@@ -16,7 +16,6 @@
# under the License.
"""Test elementwise integration."""
import numpy as np
-
import tvm
import tvm.testing
from tvm import te
@@ -62,6 +61,7 @@ def test_exp():
@tvm.testing.requires_gpu
def test_fmod():
"""Test scheduling and running fmod."""
+
# graph
def run(dtype):
size_var_n = te.size_var("n")
@@ -139,6 +139,7 @@ def test_multiple_cache_write():
schedule[cache_b0].compute_at(schedule[result_c], axis0)
schedule[result_c].bind(axis0, te.thread_axis("blockIdx.x"))
schedule[result_c].bind(axis1, te.thread_axis("threadIdx.x"))
+
# one line to build the function.
def check_device(device, host="stackvm"):
if not tvm.testing.device_enabled(host):
@@ -324,7 +325,7 @@ def try_warp_memory():
schedule[cache_read_aa].bind(axis_xi, thread_axis_tx)
@tvm.register_func("tvm_callback_cuda_compile", override=True)
- def tvm_callback_cuda_compile(code): # pylint: disable=unused-variable
+ def tvm_callback_cuda_compile(code, _): # pylint: disable=unused-variable
ptx = nvcc.compile_cuda(code)
return ptx
diff --git a/tests/python/integration/test_ewise_fpga.py
b/tests/python/integration/test_ewise_fpga.py
index 7b247d7d52..cae6364c22 100644
--- a/tests/python/integration/test_ewise_fpga.py
+++ b/tests/python/integration/test_ewise_fpga.py
@@ -18,7 +18,6 @@
import os
import numpy as np
-
import tvm
import tvm.testing
from tvm import te
@@ -28,7 +27,7 @@ os.environ["CL_CONTEXT_EMULATOR_DEVICE_INTELFPGA"] = "1"
@tvm.register_func
-def tvm_callback_vhls_postproc(code):
+def tvm_callback_vhls_postproc(code, _):
"""Hook to inspect the Vivado HLS code before actually run it"""
print(code)
return code
diff --git a/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
b/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
index 168f8c879b..5db33a1e05 100644
--- a/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
@@ -356,7 +356,7 @@ support_async = True
@tvm.register_func
-def tvm_callback_cuda_postproc(code):
+def tvm_callback_cuda_postproc(code, _):
global generated_code
global support_async
generated_code = code
