You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/06/01 23:47:21 UTC
[tvm] branch unity-staging updated (9a3bbc1d37 -> 153f70aa60)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
from 9a3bbc1d37 [MERGE] merge main into unity 2023-05-20
add 7d2579e62e [Unity][WebGPU] Try F16 support for WebGPU Backend (#14904)
add afa53ded93 [Unity][Transform] Fix scalar case in DefaultGPUSchedule (#14928)
add 6bb531fcb1 [Unity] Allow eliminating only call nodes in CSE pass (#14895)
add 17be2d71ad [Unity] Support causal mask for `R.nn.attention` (#14907)
add 4b072e05f6 [Unity][WebGPU] Fix WebGPU adapter requiring features (#14933)
add eb60d4cac1 [Unity] Cleanup Web runtime wasi (#14929)
add f15f612aa7 [Unity] Improve caching logic in webruntime (#14940)
add 286c6722e0 [Unity] Enhance web cache to add scopes (#14953)
add 66863f53f6 [Unity] Fix MergeCompositeFunctions for non-CallNode dataflow inputs (#14959)
add ec127cb0d6 [Unity][BYOC] Make CUTLASS attention rewriting aware of fp16 <-> f32 casting (#14957)
add dd2f452e4d [Unity][Transform] Fix bug for tir expression in shape in fuse_tir (#14931)
add 7d78b278d3 [Unity] support update KV cache (#14964)
add 240111fe64 [Unity][CUTLASS] Attention hot fix (#14966)
add be9f0781a9 [Unity][Frontend] Translate StableHLO to Relax (#14460)
add cb3bf4014a [Unity][Training] Enhance op gradient (#14932)
add c448c50cd1 [Unity][CUTLASS] Support more residual input shape (#14968)
add 9c0b41bc74 [Unity] Allow filtering out unwanted branches in matmul combining pass (#14971)
add a2bddcf3b2 [Unity] Add popn to kvcache (#14970)
add f796b69a3b [Unity] Fix LazyTransformParams use-def analysis and binding emission (#14974)
add 34e87ae7f8 [Unity] Fix ConvertLayout on binary elemwise ops involving scalar input (#14961)
add 44c849f144 [Unity] Reset match state when backtracking (#14984)
add 0a0b11982e [Unity][Training] Avoid problematic inputs to nll_loss in test_op_gradient_numeric (#14987)
add f360556a0f [Unity][CUTLASS] Fix circular import bug in relax cutlass backend (#15001)
add 68800fa810 [Contrib] Use f-strings for string formatting, NFC (#14893)
add 5a094bce74 [TIR] Expand unit tests for ConvertSSA (#14892)
add 53cee4bca3 [TVMScript] Round-trip DeclBuffer with undefined data pointer (#14900)
add cd4551353b [LLVM] Codegen subroutine call when CallNode::op is GlobalVar (#14901)
add dddc339d2f [Test Cases][BugFix] Add some version check to make test cases run in all PyTorch versions (#14903)
add 4f99750611 [TypoFix] fix some typo problem in keras frontend (#14916)
add d776bccc1e [Relay] add a dimension check to reject invalid input (#14925)
add 35ac2262eb [Docker] Add polly package (#14912)
add 172120a952 [TestCases] fix unreachable test cases due to outside the for-loop (#14934)
add d9c1ba60f4 [microNPU][ETHOSU] Add offloading to the NPU the nn.avg_pool2d operator with a stride > 3 (#14861)
add 1b9678cee0 [TVMScript] Allow T.target("device", host="host") to specify host (#14915)
add dbcd198670 [BUILD] Enable log before throw message in windows (#14937)
add e11913be06 [Target] Add target to all TVM callbacks (#14939)
add 1c39613811 [LLVM] Expose Host CPU Feature Detection (#14946)
add bcf7abba29 Fix pytorch axis (#14930)
add 41a616ffba [TIR] Handle subroutine calls in MakePackedAPI (#14913)
add 94f4e25a40 [TIR] Handle subroutine calls in MakeUnpackedAPI (#14914)
add 5fd49f78ad [TIR] Restrict tir.transform.CombineContextCall to host functions (#14945)
add 86ba26d854 [Bugfix][TIR] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI (#14950)
add 81056cccd5 [TIR] Preserve existing kTarget function attribute in BindTarget (#14942)
add 875217c79b [TIR] Restrict tir.transform.InstallDebugSpans to host functions (#14943)
add 3a15eafd14 [TVMScript] Prevent bool to int conversion in T.Assert condition (#14941)
add 1aeb34af33 [Codegen][LLVM] Allow void return type from PackedFunc (#14958)
add 6198c7fd8a [METAL] Fix int8 vectorized cast (#14962)
add 6eb0779442 [TIR] SplitHostDevice, handle subroutines (#14918)
add 227c4e87f3 [CI] Update the expected CI jobs list in the update_branch script (#14908)
add 4f041c96e7 [Runtime,RPC] Use f-strings for string formatting, NFC (#14967)
add c47d36e77c [Hexagon] Add support for v73, make v68 default (#14965)
add 5f1421dd0f [Bugfix][PyTorch] Support use_input_stats in instance_norm (#14963)
add 94c1b89abc [TVMScript][TIR] Parse subroutine calls with no arguments (#14919)
add 43f06ca42a [TIR] Avoid re-defining `var = arg_var` in ArgBinder (#14952)
add 4267fbf6a1 [IR][SIBuilder] (#14574)
add 7131411f0a [Bugfix][TIR][VTA] Update host-side target, even without device func (#14982)
add c8f97aa877 [TIR] Fix incorrect construction of block frames (#14993)
add c98e29bbf4 [Bugfix] Make ThreadAllReduce pass compatible with int64 (#14991)
add ea57778087 [CMAKE] Add a dummy target to defer libtvm dep (#14975)
add 4eb1a4fdf0 [Hexagon] Update instructions to compile hexagon runtime (#14948)
add 7fe58a181c [OpenCL] Improve diagnostic message (#14995)
add 443fa20e20 [RUNTIME] Update Module and Registry to use String Container (#14902)
add 9da026194f [OpenCL][Adreno] Fix conv2d when output channels < 4 (#14996)
add 8543cec133 [Hexagon] Remove "c" as aot_host_target tvm/contrib/hexagon/pytest_pl… (#14997)
add b13be936a0 [DataType] Initial support of fp8 (e4m3/e5m2) (#14863)
add 1608ca82ce [CMAKE] Introduce dummy build as an option (#15000)
add ca30b13668 [DOCKER][ADRENO] Docker setup changes for multi user environments (#15003)
add 7f02606af2 [OPENCL] Always use convert_T for type conversion (#14972)
new 153f70aa60 [MERGE] Merge main into unity 2023-06-01
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
CMakeLists.txt | 21 +-
apps/dso_plugin_module/plugin_module.cc | 3 +-
apps/hexagon_launcher/README.md | 14 +-
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 +-
ci/scripts/github/update_branch.py | 19 +-
cmake/config.cmake | 4 +-
cmake/modules/HexagonSDK.cmake | 9 +-
cmake/modules/LibInfo.cmake | 1 +
docker/install/ubuntu2204_install_llvm.sh | 2 +-
docker/install/ubuntu_install_python_package.sh | 3 +-
include/tvm/ir/si_builder.h | 103 +
include/tvm/ir/source_map.h | 46 +-
include/tvm/relax/attrs/nn.h | 3 +
include/tvm/relax/transform.h | 3 +-
include/tvm/runtime/data_type.h | 34 +
include/tvm/runtime/logging.h | 6 +-
include/tvm/runtime/module.h | 19 +-
include/tvm/runtime/packed_func.h | 2 +-
include/tvm/runtime/registry.h | 12 +-
include/tvm/runtime/relax_vm/executable.h | 6 +-
include/tvm/runtime/vm/executable.h | 4 +-
include/tvm/runtime/vm/vm.h | 2 +-
include/tvm/target/target.h | 13 +
include/tvm/tir/op.h | 3 +-
include/tvm/tir/transform.h | 52 +
jvm/core/src/test/scripts/test_add_gpu.py | 4 +-
python/gen_requirements.py | 1 +
python/tvm/_ffi/runtime_ctypes.py | 22 +
python/tvm/contrib/clang.py | 6 +-
python/tvm/contrib/cudnn.py | 6 +-
python/tvm/contrib/cutlass/attention_operation.py | 1 +
python/tvm/contrib/cutlass/build.py | 87 +-
python/tvm/contrib/cutlass/conv2d_operation.py | 36 +-
python/tvm/contrib/cutlass/gemm_operation.py | 22 +-
python/tvm/contrib/cutlass/gen_conv2d.py | 2 +-
python/tvm/contrib/cutlass/gen_gemm.py | 2 +-
python/tvm/contrib/cutlass/gen_tensor_op.py | 54 +-
python/tvm/contrib/cutlass/library.py | 23 +-
python/tvm/contrib/graph_executor.py | 2 +-
.../contrib/hexagon/profiling/process_lwp_data.py | 2 +-
python/tvm/contrib/hexagon/pytest_plugin.py | 2 +-
python/tvm/contrib/hexagon/tools.py | 8 +-
python/tvm/contrib/nvcc.py | 29 +-
python/tvm/contrib/peak.py | 4 +-
python/tvm/contrib/pickle_memoize.py | 4 +-
python/tvm/contrib/pipeline_executor.py | 20 +-
python/tvm/contrib/pipeline_executor_build.py | 22 +-
python/tvm/contrib/rocm.py | 4 +-
python/tvm/contrib/sdaccel.py | 10 +-
python/tvm/contrib/sparse.py | 8 +-
python/tvm/contrib/tar.py | 2 +-
python/tvm/contrib/target/coreml.py | 23 +-
python/tvm/contrib/target/onnx.py | 79 +-
python/tvm/contrib/target/vitis_ai.py | 2 +-
python/tvm/contrib/tf_op/module.py | 2 +-
python/tvm/contrib/utils.py | 2 +-
python/tvm/contrib/xcode.py | 6 +-
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/base.py | 17 +
python/tvm/relax/backend/contrib/__init__.py | 20 -
python/tvm/relax/backend/contrib/cutlass.py | 23 +-
python/tvm/relax/backend/patterns.py | 41 +-
python/tvm/relax/dpl/rewrite.py | 9 +-
.../tvm/relax/frontend/stablehlo}/__init__.py | 6 +-
.../frontend/stablehlo/stablehlo_translator.py | 445 +
python/tvm/relax/frontend/torch/fx_translator.py | 9 +-
python/tvm/relax/op/_op_gradient.py | 74 +-
python/tvm/relax/op/nn/nn.py | 31 +-
.../tvm/relax/transform/lazy_transform_params.py | 72 +-
python/tvm/relax/transform/legalize_ops/nn.py | 26 +-
python/tvm/relax/transform/transform.py | 21 +-
python/tvm/relay/__init__.py | 1 +
.../tvm/relay/backend/contrib/ethosu/legalize.py | 11 +-
python/tvm/relay/base.py | 2 +-
python/tvm/relay/frontend/keras.py | 4 +-
python/tvm/relay/frontend/pytorch.py | 26 +-
python/tvm/relay/op/contrib/ethosu.py | 27 +-
python/tvm/relay/op/strategy/adreno.py | 47 +-
python/tvm/rpc/base.py | 6 +-
python/tvm/rpc/client.py | 23 +-
python/tvm/rpc/minrpc.py | 4 +-
python/tvm/rpc/proxy.py | 22 +-
python/tvm/rpc/server.py | 12 +-
python/tvm/rpc/testing.py | 4 +-
python/tvm/rpc/tracker.py | 16 +-
python/tvm/runtime/container.py | 14 +-
python/tvm/runtime/executor/aot_executor.py | 2 +-
python/tvm/runtime/module.py | 38 +-
python/tvm/runtime/ndarray.py | 36 +-
python/tvm/runtime/object.py | 7 +-
python/tvm/runtime/object_generic.py | 6 +-
python/tvm/runtime/vm.py | 15 +-
python/tvm/script/ir_builder/tir/ir.py | 24 +-
python/tvm/script/parser/tir/parser.py | 14 +-
python/tvm/target/target.py | 8 +-
python/tvm/tir/op.py | 2 +-
python/tvm/tir/transform/transform.py | 84 +
src/driver/driver_api.cc | 5 +
src/ir/expr.cc | 14 +-
src/ir/si_builder.cc | 325 +
src/ir/source_map.cc | 60 +
src/relax/ir/dataflow_matcher.cc | 30 +-
src/relax/op/nn/attention.cc | 7 +-
src/relax/op/nn/attention.h | 3 +-
src/relax/op/tensor/binary.cc | 9 +
src/relax/transform/combine_parallel_matmul.cc | 41 +-
src/relax/transform/convert_layout.cc | 2 +-
src/relax/transform/eliminate_common_subexpr.cc | 21 +-
src/relax/transform/fuse_ops.cc | 3 +-
src/relax/transform/fuse_tir.cc | 3 +-
src/relax/transform/merge_composite_functions.cc | 11 +
src/relay/backend/aot_executor_codegen.cc | 2 +-
src/relay/backend/build_module.cc | 2 +-
src/relay/backend/contrib/ethosu/source_module.cc | 8 +-
src/relay/backend/graph_executor_codegen.cc | 2 +-
src/relay/backend/vm/compiler.cc | 2 +-
src/relay/backend/vm/compiler.h | 2 +-
src/relay/printer/model_library_format_printer.cc | 2 +-
src/runtime/aot_executor/aot_executor.cc | 3 +-
src/runtime/aot_executor/aot_executor.h | 2 +-
src/runtime/aot_executor/aot_executor_factory.cc | 2 +-
src/runtime/aot_executor/aot_executor_factory.h | 2 +-
src/runtime/const_loader_module.cc | 2 +-
src/runtime/contrib/coreml/coreml_runtime.h | 2 +-
src/runtime/contrib/coreml/coreml_runtime.mm | 3 +-
src/runtime/contrib/dnnl/dnnl_json_runtime.cc | 2 +-
src/runtime/contrib/ethosn/ethosn_runtime.cc | 5 +-
src/runtime/contrib/ethosn/ethosn_runtime.h | 4 +-
src/runtime/contrib/json/json_runtime.h | 4 +-
src/runtime/contrib/libtorch/libtorch_runtime.cc | 4 +-
src/runtime/contrib/onnx/onnx_module.cc | 6 +-
src/runtime/contrib/tflite/tflite_runtime.cc | 3 +-
src/runtime/contrib/tflite/tflite_runtime.h | 2 +-
src/runtime/contrib/vitis_ai/vitis_ai_runtime.cc | 3 +-
src/runtime/contrib/vitis_ai/vitis_ai_runtime.h | 2 +-
src/runtime/cuda/cuda_module.cc | 11 +-
.../cuda_graph/graph_runtime_cuda_graph.cc | 4 +-
.../graph_executor/debug/graph_executor_debug.cc | 2 +-
.../graph_executor/debug/graph_executor_debug.h | 2 +-
src/runtime/graph_executor/graph_executor.cc | 3 +-
src/runtime/graph_executor/graph_executor.h | 2 +-
.../graph_executor/graph_executor_factory.cc | 2 +-
.../graph_executor/graph_executor_factory.h | 2 +-
src/runtime/hexagon/README.md | 9 +-
src/runtime/hexagon/hexagon_module.cc | 6 +-
src/runtime/hexagon/hexagon_module.h | 6 +-
src/runtime/hexagon/rpc/simulator/session.cc | 30 +
src/runtime/library_module.cc | 5 +-
src/runtime/metadata.cc | 2 +-
src/runtime/metal/metal_module.mm | 9 +-
src/runtime/module.cc | 15 +-
src/runtime/opencl/opencl_common.h | 8 +-
src/runtime/opencl/opencl_module.cc | 14 +-
src/runtime/opencl/opencl_module_spirv.cc | 6 +-
src/runtime/opencl/sdaccel/sdaccel_module.cc | 2 +-
src/runtime/pipeline/pipeline_executor.cc | 2 +-
src/runtime/pipeline/pipeline_executor.h | 2 +-
src/runtime/registry.cc | 18 +-
src/runtime/relax_vm/executable.cc | 6 +-
src/runtime/relax_vm/lm_support.cc | 35 +
src/runtime/relax_vm/ndarray_cache_support.cc | 4 +-
src/runtime/relax_vm/vm.cc | 6 +-
src/runtime/rocm/rocm_module.cc | 9 +-
src/runtime/rpc/rpc_module.cc | 4 +-
src/runtime/stackvm/stackvm_module.cc | 9 +-
src/runtime/static_library.cc | 4 +-
src/runtime/vm/executable.cc | 6 +-
src/runtime/vm/profiler/vm.cc | 2 +-
src/runtime/vm/profiler/vm.h | 2 +-
src/runtime/vm/vm.cc | 3 +-
src/runtime/vulkan/vulkan_module.cc | 2 +-
src/runtime/vulkan/vulkan_wrapped_func.cc | 6 +-
src/runtime/vulkan/vulkan_wrapped_func.h | 6 +-
src/script/printer/tir/block.cc | 3 +-
src/script/printer/tir/buffer.cc | 29 +-
src/script/printer/tir/function.cc | 6 +-
src/script/printer/tir/stmt.cc | 24 +-
src/script/printer/tir/utils.h | 28 +-
src/support/ffi_testing.cc | 4 +-
src/support/libinfo.cc | 1 +
src/support/scalars.h | 12 +
src/target/llvm/codegen_amdgpu.cc | 10 +-
src/target/llvm/codegen_cpu.cc | 51 +-
src/target/llvm/codegen_cpu.h | 2 +-
src/target/llvm/codegen_hexagon.cc | 4 +-
src/target/llvm/codegen_llvm.cc | 110 +-
src/target/llvm/codegen_llvm.h | 61 +-
src/target/llvm/codegen_nvptx.cc | 14 +-
src/target/llvm/llvm_module.cc | 33 +-
src/target/opt/build_cuda_on.cc | 4 +-
src/target/source/codegen_aocl.cc | 2 +-
src/target/source/codegen_cuda.cc | 17 +
src/target/source/codegen_cuda.h | 5 +-
src/target/source/codegen_metal.cc | 7 +-
src/target/source/codegen_opencl.cc | 12 +-
src/target/source/codegen_vhls.cc | 5 +-
src/target/source/codegen_webgpu.cc | 28 +-
src/target/source/codegen_webgpu.h | 8 +
src/target/source/interface_c.cc | 4 +-
src/target/source/source_module.cc | 28 +-
src/target/spirv/spirv_utils.cc | 2 +-
src/target/target.cc | 15 +
src/tir/ir/stmt.cc | 3 +
src/tir/op/op.cc | 7 +
src/tir/schedule/state.cc | 2 +-
src/tir/transforms/annotate_device_regions.cc | 81 +
src/tir/transforms/combine_context_call.cc | 7 +-
src/tir/transforms/default_gpu_schedule.cc | 10 +
src/tir/transforms/dtype_conversion.cc | 101 +
src/tir/transforms/dtype_conversion.h | 165 +
src/tir/transforms/install_debug_spans.cc | 36 +-
src/tir/transforms/ir_utils.cc | 32 +-
src/tir/transforms/ir_utils.h | 13 +
src/tir/transforms/lower_device_kernel_launch.cc | 305 +
src/tir/transforms/lower_thread_allreduce.cc | 3 +-
src/tir/transforms/make_packed_api.cc | 160 +-
src/tir/transforms/make_unpacked_api.cc | 113 +-
src/tir/transforms/primfunc_utils.cc | 30 +-
src/tir/transforms/split_host_device.cc | 272 +-
...6_legalize.cc => unsupported_dtype_legalize.cc} | 264 +-
tests/cpp/si_builder_test.cc | 399 +
tests/lint/check_file_type.py | 2 +
tests/python/ci/test_ci.py | 49 +-
tests/python/contrib/test_ethosu/test_codegen.py | 32 +
tests/python/contrib/test_ethosu/test_legalize.py | 211 +
tests/python/contrib/test_hexagon/README.md | 2 +-
tests/python/frontend/keras/test_forward.py | 2 +-
tests/python/frontend/pytorch/test_forward.py | 29 +-
tests/python/integration/test_ewise.py | 5 +-
tests/python/integration/test_ewise_fpga.py | 3 +-
tests/python/relax/test_codegen_cutlass.py | 235 +-
tests/python/relax/test_dataflow_pattern.py | 112 +-
tests/python/relax/test_frontend_from_fx.py | 36 +
tests/python/relax/test_frontend_stablehlo.py | 355 +
tests/python/relax/test_op_gradient_numeric.py | 104 +-
.../test_transform_combine_parallel_matmul.py | 52 +
.../python/relax/test_transform_convert_layout.py | 45 +-
tests/python/relax/test_transform_cse.py | 32 +-
tests/python/relax/test_transform_fuse_tir.py | 76 +
tests/python/relax/test_transform_gradient.py | 6 +-
.../relax/test_transform_lazy_transform_params.py | 66 +-
.../python/relax/test_transform_legalize_ops_nn.py | 105 +-
.../test_transform_merge_composite_functions.py | 110 +
.../opencl_texture/test_conv2d_nchw_texture.py | 30 +
.../opencl_texture/test_conv2d_nhwc_texture.py | 30 +
.../test_depthwise_conv2d_nchw_texture.py | 31 +
.../test_depthwise_conv2d_nhwc_texture.py | 31 +
tests/python/tir/test_debug_info.py | 50 +-
tests/python/unittest/test_datatype_nv_fp8.py | 104 +
tests/python/unittest/test_target_codegen_llvm.py | 89 +-
tests/python/unittest/test_target_codegen_metal.py | 30 +-
.../python/unittest/test_target_codegen_opencl.py | 80 +-
.../unittest/test_target_texture_codegen_opencl.py | 38 +-
.../unittest/test_te_schedule_bound_inference.py | 56 +-
.../test_tir_transform_annotate_device_regions.py | 58 +
.../unittest/test_tir_transform_bf16_legalize.py | 5 +-
.../test_tir_transform_combine_context_call.py | 102 +-
.../unittest/test_tir_transform_convert_ssa.py | 253 +
.../test_tir_transform_device_kernel_launch.py | 193 +
.../unittest/test_tir_transform_fp8_legalize.py | 224 +
.../python/unittest/test_tir_transform_helpers.py | 112 +
.../test_tir_transform_inject_ptx_async_copy.py | 4 +-
.../python/unittest/test_tir_transform_ir_utils.py | 40 -
.../test_tir_transform_lower_warp_memory.py | 37 +-
.../unittest/test_tir_transform_make_packed_api.py | 128 +-
.../test_tir_transform_make_unpacked_api.py | 158 +-
.../test_tir_transform_split_host_device.py | 129 +-
.../unittest/test_tir_transform_thread_sync.py | 5 +-
.../test_transform_default_gpu_schedule.py | 35 +
.../python/unittest/test_tvmscript_printer_tir.py | 7 +-
tests/python/unittest/test_tvmscript_roundtrip.py | 80 +
.../python/unittest/test_tvmscript_syntax_sugar.py | 35 +
tests/scripts/setup-adreno-env.sh | 13 +-
tests/scripts/task_config_build_cpu.sh | 3 -
tests/scripts/task_config_build_i386.sh | 3 -
tests/scripts/task_python_adreno.sh | 13 +-
tests/scripts/task_python_vta_fsim.sh | 3 +
web/.gitignore | 1 +
web/Makefile | 12 +-
web/apps/node/example.js | 5 +-
web/emcc/decorate_as_wasi.py | 20 +-
web/emcc/webgpu_runtime.cc | 6 +-
web/package-lock.json | 8609 +++-----------------
web/package.json | 37 +-
web/{rollup.config.js => rollup.config.mjs} | 22 +-
web/src/compact.ts | 10 +
web/src/index.ts | 8 +-
web/src/memory.ts | 3 +-
web/src/runtime.ts | 82 +-
web/src/tvmjs_runtime_wasi.d.ts | 8 +
web/src/webgpu.ts | 10 +-
web/tests/node/test_module_load.js | 3 +-
web/tests/node/test_ndarray.js | 7 +-
web/tests/node/test_object.js | 5 +-
web/tests/node/test_packed_func.js | 7 +-
web/tests/node/test_relax_vm.js | 4 +-
web/tsconfig.json | 5 +-
web/typedoc.json | 5 +-
305 files changed, 9012 insertions(+), 9375 deletions(-)
create mode 100644 include/tvm/ir/si_builder.h
delete mode 100644 python/tvm/relax/backend/contrib/__init__.py
copy {apps/uma/_template => python/tvm/relax/frontend/stablehlo}/__init__.py (86%)
create mode 100644 python/tvm/relax/frontend/stablehlo/stablehlo_translator.py
create mode 100644 src/ir/si_builder.cc
create mode 100644 src/tir/transforms/annotate_device_regions.cc
create mode 100644 src/tir/transforms/dtype_conversion.cc
create mode 100644 src/tir/transforms/dtype_conversion.h
create mode 100644 src/tir/transforms/lower_device_kernel_launch.cc
rename src/tir/transforms/{bf16_legalize.cc => unsupported_dtype_legalize.cc} (69%)
create mode 100644 tests/cpp/si_builder_test.cc
create mode 100644 tests/python/relax/test_frontend_stablehlo.py
create mode 100644 tests/python/unittest/test_datatype_nv_fp8.py
create mode 100644 tests/python/unittest/test_tir_transform_annotate_device_regions.py
create mode 100644 tests/python/unittest/test_tir_transform_convert_ssa.py
create mode 100644 tests/python/unittest/test_tir_transform_device_kernel_launch.py
create mode 100644 tests/python/unittest/test_tir_transform_fp8_legalize.py
delete mode 100644 tests/python/unittest/test_tir_transform_ir_utils.py
rename web/{rollup.config.js => rollup.config.mjs} (71%)
create mode 100644 web/src/tvmjs_runtime_wasi.d.ts
[tvm] 01/01: [MERGE] Merge main into unity 2023-06-01
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 153f70aa60b82c4f7670ee32dceee859a431fdd0
Merge: f360556a0f 7f02606af2
Author: tqchen <ti...@gmail.com>
AuthorDate: Thu Jun 1 19:46:25 2023 -0400
[MERGE] Merge main into unity 2023-06-01
Updated the runtime to latest
cutlass conflict resolutions
CMakeLists.txt | 21 +-
apps/dso_plugin_module/plugin_module.cc | 3 +-
apps/hexagon_launcher/README.md | 14 +-
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 +-
ci/scripts/github/update_branch.py | 19 +-
cmake/config.cmake | 4 +-
cmake/modules/HexagonSDK.cmake | 9 +-
cmake/modules/LibInfo.cmake | 1 +
docker/install/ubuntu2204_install_llvm.sh | 2 +-
docker/install/ubuntu_install_python_package.sh | 3 +-
include/tvm/ir/si_builder.h | 103 ++++++
include/tvm/ir/source_map.h | 46 ++-
include/tvm/runtime/data_type.h | 34 ++
include/tvm/runtime/logging.h | 6 +-
include/tvm/runtime/module.h | 19 +-
include/tvm/runtime/packed_func.h | 2 +-
include/tvm/runtime/registry.h | 12 +-
include/tvm/runtime/relax_vm/executable.h | 6 +-
include/tvm/runtime/vm/executable.h | 4 +-
include/tvm/runtime/vm/vm.h | 2 +-
include/tvm/target/target.h | 13 +
include/tvm/tir/op.h | 3 +-
include/tvm/tir/transform.h | 52 +++
jvm/core/src/test/scripts/test_add_gpu.py | 4 +-
python/gen_requirements.py | 1 +
python/tvm/_ffi/runtime_ctypes.py | 22 ++
python/tvm/contrib/clang.py | 6 +-
python/tvm/contrib/cudnn.py | 6 +-
python/tvm/contrib/cutlass/build.py | 29 +-
python/tvm/contrib/cutlass/conv2d_operation.py | 11 +-
python/tvm/contrib/cutlass/gemm_operation.py | 22 +-
python/tvm/contrib/cutlass/gen_conv2d.py | 2 +-
python/tvm/contrib/cutlass/gen_gemm.py | 2 +-
python/tvm/contrib/cutlass/gen_tensor_op.py | 49 +--
python/tvm/contrib/cutlass/library.py | 23 +-
python/tvm/contrib/graph_executor.py | 2 +-
.../contrib/hexagon/profiling/process_lwp_data.py | 2 +-
python/tvm/contrib/hexagon/pytest_plugin.py | 2 +-
python/tvm/contrib/hexagon/tools.py | 8 +-
python/tvm/contrib/nvcc.py | 29 +-
python/tvm/contrib/peak.py | 4 +-
python/tvm/contrib/pickle_memoize.py | 4 +-
python/tvm/contrib/pipeline_executor.py | 20 +-
python/tvm/contrib/pipeline_executor_build.py | 22 +-
python/tvm/contrib/rocm.py | 4 +-
python/tvm/contrib/sdaccel.py | 10 +-
python/tvm/contrib/sparse.py | 8 +-
python/tvm/contrib/tar.py | 2 +-
python/tvm/contrib/target/coreml.py | 23 +-
python/tvm/contrib/target/onnx.py | 79 ++--
python/tvm/contrib/target/vitis_ai.py | 2 +-
python/tvm/contrib/tf_op/module.py | 2 +-
python/tvm/contrib/utils.py | 2 +-
python/tvm/contrib/xcode.py | 6 +-
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/base.py | 17 +
python/tvm/relay/__init__.py | 1 +
.../tvm/relay/backend/contrib/ethosu/legalize.py | 11 +-
python/tvm/relay/base.py | 2 +-
python/tvm/relay/frontend/keras.py | 4 +-
python/tvm/relay/frontend/pytorch.py | 26 +-
python/tvm/relay/op/contrib/ethosu.py | 27 +-
python/tvm/relay/op/strategy/adreno.py | 47 ++-
python/tvm/rpc/base.py | 6 +-
python/tvm/rpc/client.py | 23 +-
python/tvm/rpc/minrpc.py | 4 +-
python/tvm/rpc/proxy.py | 22 +-
python/tvm/rpc/server.py | 12 +-
python/tvm/rpc/testing.py | 4 +-
python/tvm/rpc/tracker.py | 16 +-
python/tvm/runtime/container.py | 14 +-
python/tvm/runtime/executor/aot_executor.py | 2 +-
python/tvm/runtime/module.py | 38 +-
python/tvm/runtime/ndarray.py | 36 +-
python/tvm/runtime/object.py | 7 +-
python/tvm/runtime/object_generic.py | 6 +-
python/tvm/runtime/vm.py | 15 +-
python/tvm/script/ir_builder/tir/ir.py | 24 +-
python/tvm/script/parser/tir/parser.py | 14 +-
python/tvm/target/target.py | 8 +-
python/tvm/tir/op.py | 2 +-
python/tvm/tir/transform/transform.py | 84 +++++
src/driver/driver_api.cc | 5 +
src/ir/expr.cc | 14 +-
src/ir/si_builder.cc | 325 +++++++++++++++++
src/ir/source_map.cc | 60 ++++
src/relay/backend/aot_executor_codegen.cc | 2 +-
src/relay/backend/build_module.cc | 2 +-
src/relay/backend/contrib/ethosu/source_module.cc | 8 +-
src/relay/backend/graph_executor_codegen.cc | 2 +-
src/relay/backend/vm/compiler.cc | 2 +-
src/relay/backend/vm/compiler.h | 2 +-
src/relay/printer/model_library_format_printer.cc | 2 +-
src/runtime/aot_executor/aot_executor.cc | 3 +-
src/runtime/aot_executor/aot_executor.h | 2 +-
src/runtime/aot_executor/aot_executor_factory.cc | 2 +-
src/runtime/aot_executor/aot_executor_factory.h | 2 +-
src/runtime/const_loader_module.cc | 2 +-
src/runtime/contrib/coreml/coreml_runtime.h | 2 +-
src/runtime/contrib/coreml/coreml_runtime.mm | 3 +-
src/runtime/contrib/dnnl/dnnl_json_runtime.cc | 2 +-
src/runtime/contrib/ethosn/ethosn_runtime.cc | 5 +-
src/runtime/contrib/ethosn/ethosn_runtime.h | 4 +-
src/runtime/contrib/json/json_runtime.h | 4 +-
src/runtime/contrib/libtorch/libtorch_runtime.cc | 4 +-
src/runtime/contrib/onnx/onnx_module.cc | 6 +-
src/runtime/contrib/tflite/tflite_runtime.cc | 3 +-
src/runtime/contrib/tflite/tflite_runtime.h | 2 +-
src/runtime/contrib/vitis_ai/vitis_ai_runtime.cc | 3 +-
src/runtime/contrib/vitis_ai/vitis_ai_runtime.h | 2 +-
src/runtime/cuda/cuda_module.cc | 11 +-
.../cuda_graph/graph_runtime_cuda_graph.cc | 4 +-
.../graph_executor/debug/graph_executor_debug.cc | 2 +-
.../graph_executor/debug/graph_executor_debug.h | 2 +-
src/runtime/graph_executor/graph_executor.cc | 3 +-
src/runtime/graph_executor/graph_executor.h | 2 +-
.../graph_executor/graph_executor_factory.cc | 2 +-
.../graph_executor/graph_executor_factory.h | 2 +-
src/runtime/hexagon/README.md | 9 +-
src/runtime/hexagon/hexagon_module.cc | 6 +-
src/runtime/hexagon/hexagon_module.h | 6 +-
src/runtime/hexagon/rpc/simulator/session.cc | 30 ++
src/runtime/library_module.cc | 5 +-
src/runtime/metadata.cc | 2 +-
src/runtime/metal/metal_module.mm | 9 +-
src/runtime/module.cc | 15 +-
src/runtime/opencl/opencl_common.h | 8 +-
src/runtime/opencl/opencl_module.cc | 14 +-
src/runtime/opencl/opencl_module_spirv.cc | 6 +-
src/runtime/opencl/sdaccel/sdaccel_module.cc | 2 +-
src/runtime/pipeline/pipeline_executor.cc | 2 +-
src/runtime/pipeline/pipeline_executor.h | 2 +-
src/runtime/registry.cc | 18 +-
src/runtime/relax_vm/executable.cc | 6 +-
src/runtime/relax_vm/ndarray_cache_support.cc | 4 +-
src/runtime/relax_vm/vm.cc | 6 +-
src/runtime/rocm/rocm_module.cc | 9 +-
src/runtime/rpc/rpc_module.cc | 4 +-
src/runtime/stackvm/stackvm_module.cc | 9 +-
src/runtime/static_library.cc | 4 +-
src/runtime/vm/executable.cc | 6 +-
src/runtime/vm/profiler/vm.cc | 2 +-
src/runtime/vm/profiler/vm.h | 2 +-
src/runtime/vm/vm.cc | 3 +-
src/runtime/vulkan/vulkan_module.cc | 2 +-
src/runtime/vulkan/vulkan_wrapped_func.cc | 6 +-
src/runtime/vulkan/vulkan_wrapped_func.h | 6 +-
src/script/printer/tir/block.cc | 3 +-
src/script/printer/tir/buffer.cc | 29 +-
src/script/printer/tir/function.cc | 6 +-
src/script/printer/tir/stmt.cc | 24 +-
src/script/printer/tir/utils.h | 28 +-
src/support/ffi_testing.cc | 4 +-
src/support/libinfo.cc | 1 +
src/support/scalars.h | 12 +
src/target/llvm/codegen_amdgpu.cc | 10 +-
src/target/llvm/codegen_cpu.cc | 51 +--
src/target/llvm/codegen_cpu.h | 2 +-
src/target/llvm/codegen_hexagon.cc | 4 +-
src/target/llvm/codegen_llvm.cc | 110 +++++-
src/target/llvm/codegen_llvm.h | 61 +++-
src/target/llvm/codegen_nvptx.cc | 14 +-
src/target/llvm/llvm_module.cc | 33 +-
src/target/opt/build_cuda_on.cc | 4 +-
src/target/source/codegen_aocl.cc | 2 +-
src/target/source/codegen_cuda.cc | 17 +
src/target/source/codegen_cuda.h | 5 +-
src/target/source/codegen_metal.cc | 7 +-
src/target/source/codegen_opencl.cc | 12 +-
src/target/source/codegen_vhls.cc | 5 +-
src/target/source/codegen_webgpu.cc | 8 +-
src/target/source/interface_c.cc | 4 +-
src/target/source/source_module.cc | 28 +-
src/target/spirv/spirv_utils.cc | 2 +-
src/target/target.cc | 15 +
src/tir/ir/stmt.cc | 3 +
src/tir/op/op.cc | 7 +
src/tir/schedule/state.cc | 2 +-
src/tir/transforms/annotate_device_regions.cc | 81 +++++
src/tir/transforms/combine_context_call.cc | 7 +-
src/tir/transforms/dtype_conversion.cc | 101 ++++++
src/tir/transforms/dtype_conversion.h | 165 +++++++++
src/tir/transforms/install_debug_spans.cc | 36 +-
src/tir/transforms/ir_utils.cc | 32 +-
src/tir/transforms/ir_utils.h | 13 +
src/tir/transforms/lower_device_kernel_launch.cc | 305 ++++++++++++++++
src/tir/transforms/lower_thread_allreduce.cc | 3 +-
src/tir/transforms/make_packed_api.cc | 160 ++++++---
src/tir/transforms/make_unpacked_api.cc | 113 +++++-
src/tir/transforms/primfunc_utils.cc | 30 +-
src/tir/transforms/split_host_device.cc | 272 ++++----------
...6_legalize.cc => unsupported_dtype_legalize.cc} | 264 +++++++++-----
tests/cpp/si_builder_test.cc | 399 +++++++++++++++++++++
tests/python/ci/test_ci.py | 49 +--
tests/python/contrib/test_ethosu/test_codegen.py | 32 ++
tests/python/contrib/test_ethosu/test_legalize.py | 211 +++++++++++
tests/python/contrib/test_hexagon/README.md | 2 +-
tests/python/frontend/keras/test_forward.py | 2 +-
tests/python/frontend/pytorch/test_forward.py | 29 +-
tests/python/integration/test_ewise.py | 5 +-
tests/python/integration/test_ewise_fpga.py | 3 +-
.../opencl_texture/test_conv2d_nchw_texture.py | 30 ++
.../opencl_texture/test_conv2d_nhwc_texture.py | 30 ++
.../test_depthwise_conv2d_nchw_texture.py | 31 ++
.../test_depthwise_conv2d_nhwc_texture.py | 31 ++
tests/python/tir/test_debug_info.py | 50 ++-
tests/python/unittest/test_datatype_nv_fp8.py | 104 ++++++
tests/python/unittest/test_target_codegen_llvm.py | 89 ++++-
tests/python/unittest/test_target_codegen_metal.py | 30 +-
.../python/unittest/test_target_codegen_opencl.py | 80 ++---
.../unittest/test_target_texture_codegen_opencl.py | 38 +-
.../unittest/test_te_schedule_bound_inference.py | 56 ++-
.../test_tir_transform_annotate_device_regions.py | 58 +++
.../unittest/test_tir_transform_bf16_legalize.py | 5 +-
.../test_tir_transform_combine_context_call.py | 102 ++++--
.../unittest/test_tir_transform_convert_ssa.py | 253 +++++++++++++
.../test_tir_transform_device_kernel_launch.py | 193 ++++++++++
.../unittest/test_tir_transform_fp8_legalize.py | 224 ++++++++++++
.../python/unittest/test_tir_transform_helpers.py | 112 ++++++
.../test_tir_transform_inject_ptx_async_copy.py | 4 +-
.../python/unittest/test_tir_transform_ir_utils.py | 40 ---
.../test_tir_transform_lower_warp_memory.py | 37 +-
.../unittest/test_tir_transform_make_packed_api.py | 128 ++++++-
.../test_tir_transform_make_unpacked_api.py | 158 +++++++-
.../test_tir_transform_split_host_device.py | 129 ++++++-
.../unittest/test_tir_transform_thread_sync.py | 5 +-
.../python/unittest/test_tvmscript_printer_tir.py | 7 +-
tests/python/unittest/test_tvmscript_roundtrip.py | 80 +++++
.../python/unittest/test_tvmscript_syntax_sugar.py | 35 ++
tests/scripts/setup-adreno-env.sh | 13 +-
tests/scripts/task_config_build_cpu.sh | 3 -
tests/scripts/task_config_build_i386.sh | 3 -
tests/scripts/task_python_adreno.sh | 13 +-
tests/scripts/task_python_vta_fsim.sh | 3 +
web/emcc/webgpu_runtime.cc | 6 +-
242 files changed, 5684 insertions(+), 1394 deletions(-)
diff --cc include/tvm/runtime/module.h
index deb3deae42,60e3535319..9b67b12579
--- a/include/tvm/runtime/module.h
+++ b/include/tvm/runtime/module.h
@@@ -217,11 -216,7 +216,11 @@@ class TVM_DLL ModuleNode : public Objec
* \param name name of the function.
* \return The corresponding function.
*/
- const PackedFunc* GetFuncFromEnv(const std::string& name);
+ const PackedFunc* GetFuncFromEnv(const String& name);
+
+ /*! \brief Clear all imports of the module. */
+ void ClearImports() { imports_.clear(); }
+
/*! \return The module it imports from */
const std::vector<Module>& imports() const { return imports_; }
diff --cc include/tvm/runtime/relax_vm/executable.h
index 5833cb4718,0000000000..d3bf9d356e
mode 100644,000000..100644
--- a/include/tvm/runtime/relax_vm/executable.h
+++ b/include/tvm/runtime/relax_vm/executable.h
@@@ -1,213 -1,0 +1,213 @@@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/runtime/relax_vm/executable.h
+ */
+#ifndef TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
+#define TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
+
+#include <tvm/runtime/container/closure.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "./bytecode.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+/*!
+ * \brief Information entry in executable function table.
+ *
+ * Contains metadata about the compiled function, as
+ * well as the compiled VM instructions.
+ */
+struct VMFuncInfo {
+ /*! \brief kind of the function. */
+ enum class FuncKind : int {
+ /*! \brief system level packed function */
+ kPackedFunc = 0,
+ /*! \brief VM function. */
+ kVMFunc = 1,
+ /*! \brief VMTIR function. */
+ kVMTIRFunc = 2,
+ };
+ /*! \brief The kind of function. */
+ FuncKind kind;
+ /*! \brief The function's name, global symbol */
+ std::string name;
+ /*! \brief The start instruction index of the function. */
+ Index start_instr = 0;
+ /*! \brief The end instruction index of the function. */
+ Index end_instr = 0;
+ /*! \brief The number of arguments of the function. */
+ Index num_args = 0;
+ /*! \brief The register file size of the function. */
+ Index register_file_size = 0;
+ /*! \brief The function parameter names.*/
+ std::vector<std::string> param_names;
+
+ // defined customized loader save
+ void Save(dmlc::Stream* writer) const;
+ bool Load(dmlc::Stream* reader);
+};
+
+/*!
+ * \brief The executable emitted by the VM compiler.
+ *
+ * The executable contains information (e.g. data in different memory regions)
+ * to run in a virtual machine.
+ */
+class Executable : public runtime::ModuleNode {
+ public:
+ /*!
+ * \brief Get a PackedFunc from the executable module.
+ * \param name the name of the function.
+ * \param sptr_to_self The shared_ptr that points to this module node.
+ * \return PackedFunc or nullptr when it is not available.
+ */
- PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
++ PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) final;
+
+ /*! \brief Get the property of the runtime module .*/
+ int GetPropertyMask() const final { return ModulePropertyMask::kBinarySerializable; };
+
+ /*!
+ * \brief Print the detailed statistics of the given code, i.e. number of
+ * globals and constants, etc.
+ * \return The statistics represented by a string.
+ */
+ std::string Stats() const;
+ /*!
+ * \brief Get the i-th instruction from the executable.
+ * \param i The index of the instruction to be fetched.
+ * \return The instruction.
+ */
+ Instruction GetInstruction(Index i) const;
+ /*!
+ * \brief Set j-th byte data of i-th instruction to val.
+ * \param i The index of the instruction to be updated.
+ * \param j The index of the byte data of the instruction to be updated.
+ * \param val The value to be set
+ */
+ void SetInstructionData(Index i, Index j, ExecWord val);
+ /*!
+ * \brief Print the instructions as text format.
+ * \return The text format of the instructions.
+ */
+ String AsText() const;
+ /*!
+ * \brief Print the instructions as python program.
+ * \return The python program of the instructions, represented by a string.
+ */
+ String AsPython() const;
+ /*!
+ * \brief Write the Executable to the binary stream in serialized form.
+ * \param stream The binary stream to save the executable to.
+ */
+ void SaveToBinary(dmlc::Stream* stream) final;
+ /*!
+ * \brief Load Executable from the binary stream in serialized form.
+ * \param stream The binary stream that load the executable from.
+ * \return The loaded executable, in the form of a `runtime::Module`.
+ */
+ static Module LoadFromBinary(void* stream);
+ /*!
+ * \brief Write the Executable to the provided path as a file containing its serialized content.
+ * \param file_name The name of the file to write the serialized data to.
+ * \param format The target format of the saved file.
+ */
- void SaveToFile(const std::string& file_name, const std::string& format) final;
++ void SaveToFile(const String& file_name, const String& format) final;
+ /*!
+ * \brief Load Executable from the file.
+ * \param file_name The path of the file that load the executable from.
+ * \return The loaded executable, in the form of a `runtime::Module`.
+ */
- static Module LoadFromFile(const std::string& file_name);
++ static Module LoadFromFile(const String& file_name);
+
+ /*! \brief The virtual machine's function table. */
+ std::vector<VMFuncInfo> func_table;
+ /*! \brief A map from globals (as strings) to their index in the function map. */
+ std::unordered_map<std::string, Index> func_map;
+ /*! \brief The global constant pool. */
+ std::vector<TVMRetValue> constants;
+ /*! \brief The offset of instruction. */
+ std::vector<Index> instr_offset;
+ /*! \brief The byte data of instruction. */
+ std::vector<ExecWord> instr_data;
+
+ virtual ~Executable() {}
+
+ const char* type_key() const final { return "relax.Executable"; }
+
+ private:
+ /*!
+ * \brief Save the globals.
+ * \param strm The input stream.
+ */
+ void SaveGlobalSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the constant pool.
+ * \param strm The input stream.
+ */
+ void SaveConstantSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the instructions.
+ * \param strm The input stream.
+ */
+ void SaveCodeSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the packed functions.
+ * \param strm The input stream.
+ */
+ void SavePackedFuncNames(dmlc::Stream* strm);
+ /*!
+ * \brief Load the globals.
+ * \param strm The input stream.
+ */
+ void LoadGlobalSection(dmlc::Stream* strm);
+ /*!
+ * \brief Load the constant pool.
+ * \param strm The input stream.
+ */
+ void LoadConstantSection(dmlc::Stream* strm);
+ /*!
+ * \brief Load the instructions.
+ * \param strm The input stream.
+ */
+ void LoadCodeSection(dmlc::Stream* strm);
+ /*!
+ * \brief Save the packed functions.
+ * \param strm The input stream.
+ */
+ void LoadPackedFuncNames(dmlc::Stream* strm);
+};
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
+
+namespace dmlc {
+DMLC_DECLARE_TRAITS(has_saveload, ::tvm::runtime::relax_vm::VMFuncInfo, true);
+} // namespace dmlc
+#endif // TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
diff --cc python/tvm/contrib/cutlass/build.py
index 07583a4851,0fddee54e0..f25d5eeb95
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@@ -72,9 -63,8 +70,9 @@@ def _get_cutlass_compile_options(sm, th
"-Xcompiler=-fno-strict-aliasing",
"-O3",
"-std=c++17",
- "-I" + cutlass_include,
- "-I" + cutlass_util_include,
- "-I" + cutlass_attention_include,
+ f"-I{cutlass_include}",
+ f"-I{cutlass_util_include}",
++ f"-I{cutlass_attention_include}",
]
if use_fast_math:
kwargs["options"].append("-DCUTLASS_USE_TANH_FOR_SIGMOID")
diff --cc python/tvm/contrib/cutlass/gen_conv2d.py
index 9e9e16426b,3887fc2e2e..3a9cbf1e84
--- a/python/tvm/contrib/cutlass/gen_conv2d.py
+++ b/python/tvm/contrib/cutlass/gen_conv2d.py
@@@ -185,13 -179,9 +185,13 @@@ class CutlassConv2DProfiler
def __init__(self, sm, cutlass_path, binary_path):
self.gemm_profiler = CutlassGemmProfiler(sm, cutlass_path, binary_path)
self.sm = sm
- assert sm in GENERATOR_FUNC_TABLE, "sm%d not supported yet." % sm
+ assert sm in GENERATOR_FUNC_TABLE, f"sm{sm} not supported yet."
self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
- self.cache = {}
+ self.cache_path = os.path.join(binary_path, "cutlass_conv2d_cache.pickle")
+ if os.path.exists(self.cache_path):
+ self.cache = pickle.load(open(self.cache_path, "rb"))
+ else:
+ self.cache = {}
def get_default(
self,
diff --cc python/tvm/contrib/cutlass/gen_gemm.py
index 0ea6231b81,78e19f510d..a3204791f0
--- a/python/tvm/contrib/cutlass/gen_gemm.py
+++ b/python/tvm/contrib/cutlass/gen_gemm.py
@@@ -194,24 -149,13 +194,24 @@@ class CutlassGemmProfiler
"""Profile all candidate kernels and select the best one."""
def __init__(self, sm, cutlass_path, binary_path):
- assert sm in GENERATOR_FUNC_TABLE and sm in DEFAULT_KERNELS, "sm%d not supported yet." % sm
+ assert sm in GENERATOR_FUNC_TABLE and sm in DEFAULT_KERNELS, f"sm{sm} not supported yet."
self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
self.sm = sm
- self.cache = {}
+ self.cache_path = os.path.join(binary_path, "cutlass_gemm_cache.pickle")
+ if os.path.exists(self.cache_path):
+ self.cache = pickle.load(open(self.cache_path, "rb"))
+ else:
+ self.cache = {}
def get_default(
- self, op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32=True, batched=False
+ self,
+ op_type,
+ out_dtype,
+ arg0_dtype,
+ arg1_dtype,
+ use_3xtf32=True,
+ batched=False,
+ layout_b=LayoutType.ColumnMajor,
):
"""Return the default kernel for the requested architecture.
For now, the default kernel was picked arbitrary.
diff --cc python/tvm/contrib/cutlass/gen_tensor_op.py
index 0caf62043e,855d8dc2d1..be3bb289cf
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@@ -96,14 -85,11 +90,12 @@@ def generate_sm50_simt(out_dtype, arg0_
DataType.f32,
DataType.f32,
DataType.f32,
+ DataType.f32,
OpcodeClass.Simt,
MathOperation.multiply_add,
- ),
- ]
- alignment_constraints = [
- 1,
+ )
]
+ alignment_constraints = [1]
tile_descriptions = [
([128, 128, 8], 2, [4, 2, 1], min_cc, max_cc),
([128, 64, 8], 2, [2, 2, 1], min_cc, max_cc),
@@@ -687,100 -575,9 +668,100 @@@ def instantiate_template(func_name, ann
attrs["split_k_slices"] = str(re.search(r"splitk(\d+)", op_name).group(1))
else:
attrs["split_k_mode"] = "kSerial"
- attrs["split_k_slices"] = "1"
+ attrs["split_k_slices"] = 1
+
+ if "residual_shape" in annotations:
+ attrs["residual_shape"] = annotations["residual_shape"]
+
+ code = instantiate_conv2d_template(attrs)
+ return CodegenResult(code, headers)
- code = instantiate_conv2d_template(attrs, func_args)
+ elif "attention" in func_name:
+ headers.append("kernel_forward.h")
+ data_type = dtype_map[annotations["arg0_dtype"]]
+ attrs["data_type"] = DataTypeTag[data_type]
+ attrs["num_batches"] = b = annotations["num_batches"]
+ attrs["num_queries"] = s = annotations["num_queries"]
+ attrs["num_keys"] = annotations["num_keys"]
+ attrs["num_heads"] = n = annotations["num_heads"]
+ attrs["head_dim"] = h = annotations["head_dim"]
+ attrs["head_dim_value"] = h_v = annotations["head_dim_value"]
+ data_type_size = DataTypeSize[data_type]
+ if (data_type_size * h // 8) % 16 == 0 and (data_type_size * h_v // 8) % 16 == 0:
+ attrs["kIsAligned"] = True
+ elif (h % 4 == 0) and (h_v % 4 == 0):
+ attrs["kIsAligned"] = False
+ else:
+ raise NotImplementedError()
+ if h_v > 64:
+ attrs["kQueriesPerBlock"] = 32
+ attrs["kKeysPerBlock"] = 128
+ attrs["kSingleValueIteration"] = h_v <= 128
+ else:
+ attrs["kQueriesPerBlock"] = 64
+ attrs["kKeysPerBlock"] = 64
+ attrs["kSingleValueIteration"] = True
+ attrs["output_size"] = b * s * n * h_v
+ attrs["scale"] = (
+ float(1 / math.sqrt(h.value)) if annotations["scale"] is None else annotations["scale"]
+ )
+ attrs["custom_mask_type"] = annotations["custom_mask_type"]
+
+ assert (
+ attrs["scale"] > 0 or attrs["scale"] < 0
+ ), "Cutlass may generate nan occasionally when scale == 0.0"
+ attrs["arch"] = "cutlass::arch::Sm{}".format(annotations["arch"])
+ attrs["kSupportsDropout"] = False
+ attrs["qkv_layout"] = annotations["qkv_layout"]
+
+ for arg in func_args:
+ if "workspace" in arg:
+ attrs["workspace"] = arg
+
+ if attrs["qkv_layout"] == "default":
+ attrs["query"] = func_args[0]
+ attrs["key"] = func_args[1]
+ attrs["value"] = func_args[2]
+ if len(func_args) > 4: # +1 for workspace, the last arg
+ attrs["bias"] = func_args[3]
+ elif attrs["qkv_layout"] == "qkv_stacked":
+ attrs["qkv"] = func_args[0]
+ if len(func_args) > 5: # +1 for workspace, the last arg
+ attrs["bias"] = func_args[4]
+ else:
+ raise NotImplementedError()
+ if "bias" in attrs:
+ attrs["kSupportsBias"] = True
+ if len(annotations["bias_shape"]) == 4:
+ strides = "p.num_keys"
+ if annotations["bias_shape"][2] == 1:
+ attrs["bias_strideM"] = 0
+ else:
+ attrs["bias_strideM"] = strides
+ strides = f"p.num_queries * {strides}"
+ if annotations["bias_shape"][1] == 1:
+ attrs["bias_strideH"] = 0
+ else:
+ attrs["bias_strideH"] = strides
+ strides = f"p.num_heads * {strides}"
+ if annotations["bias_shape"][0] == 1:
+ attrs["bias_strideB"] = 0
+ else:
+ attrs["bias_strideB"] = strides
+ else:
+ raise NotImplementedError()
+ else:
+ # To support negative scale in current Cutlass implementation,
+ # kSupportsBias should be set true, or there are nan's as result.
+ attrs["kSupportsBias"] = attrs["scale"] < 0
+ code = instantiate_attention_template(attrs)
+ return CodegenResult(code, headers)
+ elif "layer_norm" in func_name:
+ headers.append("cutlass/util/device_layernorm.h")
+ headers.append("cutlass/layout/matrix.h")
+ attrs = {"input": func_args[0], "gamma": func_args[1], "beta": func_args[2]}
+ attrs.update(dict(annotations))
+ code = instantiate_layer_norm_template(attrs)
return CodegenResult(code, headers)
- raise ValueError("Do not have a template for {}".format(func_name))
+ raise ValueError(f"Do not have a template for {func_name}")
diff --cc python/tvm/contrib/cutlass/library.py
index ead5804b59,dff166fb58..2b8bb830e1
--- a/python/tvm/contrib/cutlass/library.py
+++ b/python/tvm/contrib/cutlass/library.py
@@@ -145,13 -139,7 +141,13 @@@ def substitute_template(template, value
while changed:
changed = False
for key, value in values.items():
+ if isinstance(value, (int, IntImm)):
+ value = str(int(value))
+ if isinstance(value, (float, FloatImm)):
+ value = str(float(value))
+ elif isinstance(value, bool):
+ value = str(value).lower()
- regex = "\\$\\{%s\\}" % key
+ regex = f"\\$\\{{{key}\\}}"
newtext = re.sub(regex, value, text)
if newtext != text:
changed = True
diff --cc python/tvm/rpc/proxy.py
index a6709a3421,687d47264b..b14b02eb56
--- a/python/tvm/rpc/proxy.py
+++ b/python/tvm/rpc/proxy.py
@@@ -267,14 -252,9 +265,14 @@@ class ProxyServerHandler(object)
)
logging.info("Serving RPC index html page at http://localhost:%d", web_port)
resource_files = resource_files if resource_files else []
- for fname in resource_files:
+ for item in resource_files:
+ prefix, fname = item
+ if not prefix.endswith("/"):
+ prefix += "/"
+ if not prefix.startswith("/"):
+ prefix = "/" + prefix
basename = os.path.basename(fname)
- pair = (r"%s%s" % (prefix, basename), RequestHandler, {"file_path": fname})
- pair = (rf"/{basename}", RequestHandler, {"file_path": fname})
++ pair = (rf"{prefix}{basename}", RequestHandler, {"file_path": fname})
handlers.append(pair)
logging.info(pair)
self.app = tornado.web.Application(handlers)
diff --cc src/runtime/relax_vm/executable.cc
index 2090a3b254,0000000000..5abed9bbbb
mode 100644,000000..100644
--- a/src/runtime/relax_vm/executable.cc
+++ b/src/runtime/relax_vm/executable.cc
@@@ -1,583 -1,0 +1,583 @@@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/runtime/relax_vm/executable.cc
+ */
+
+#include <dmlc/memory_io.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/relax_vm/executable.h>
+#include <tvm/runtime/relax_vm/vm.h>
+
+#include <functional>
+#include <sstream>
+
+#include "../file_utils.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+/*! \brief The magic number for the serialized VM bytecode file */
+constexpr uint64_t kTVMVMBytecodeMagic = 0xD225DE2F4214151D;
+
+/*! \brief Possible types in the constant pool */
+enum ConstantType : int {
+ kNDArray = 0,
+ kDLDataType = 1,
+ kShapeTuple = 2,
+ kString = 3,
+ kInt = 4,
+};
+
+#define STREAM_CHECK(val, section) \
+ ICHECK(val) << "Invalid VM file format in the " << section << " section." \
+ << "\n";
+
- PackedFunc Executable::GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
++PackedFunc Executable::GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "stats") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->Stats(); });
+ } else if (name == "as_text") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->AsText(); });
+ } else if (name == "as_python") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->AsPython(); });
+ } else if (name == "vm_load_executable") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ObjectPtr<VirtualMachine> vm = VirtualMachine::Create();
+ ICHECK(sptr_to_self.get() == this);
+ vm->LoadExecutable(GetObjectPtr<Executable>(this));
+ *rv = Module(vm);
+ });
+ } else if (name == "vm_profiler_load_executable") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ObjectPtr<VirtualMachine> vm = VirtualMachine::CreateProfiler();
+ ICHECK(sptr_to_self.get() == this);
+ vm->LoadExecutable(GetObjectPtr<Executable>(this));
+ *rv = Module(vm);
+ });
+ }
+ return nullptr;
+}
+
+std::string Executable::Stats() const {
+ std::ostringstream oss;
+ oss << "Relax VM executable statistics:" << std::endl;
+
+ // Get the number of constants.
+ // If the constant is an NDArray, get the shape of each of them.
+ // If the constant is an DLDataType, get the data type of each of them.
+ oss << " Constant pool (# " << constants.size() << "): [";
+ for (const auto& it : constants) {
+ if (it.IsObjectRef<runtime::NDArray>()) {
+ const auto ndarray = it.operator tvm::runtime::NDArray();
+ const auto& shape = ndarray.Shape();
+ // Scalar
+ if (shape.empty()) {
+ oss << "scalar, ";
+ continue;
+ }
+ oss << "[";
+ for (auto s : shape) {
+ oss << s << ", ";
+ }
+ oss.seekp(-2, oss.cur);
+ oss << "], ";
+ } else if (it.IsObjectRef<ShapeTuple>()) {
+ ShapeTuple shape = it.operator ShapeTuple();
+ oss << "shapetuple[";
+ for (size_t i = 0; i < shape.size(); ++i) {
+ oss << shape.at(i) << ", ";
+ }
+ oss.seekp(-2, oss.cur);
+ oss << "], ";
+ } else if (it.IsObjectRef<String>()) {
+ std::string f = it.AsObjectRef<tvm::runtime::String>().operator std::string();
+ oss << "\"";
+ oss << f;
+ oss << "\", ";
+ } else if (it.type_code() == kDLInt) {
+ oss << static_cast<int64_t>(it);
+ oss << ", ";
+ } else {
+ try {
+ DLDataType dtype = it.operator DLDataType();
+ oss << dtype;
+ oss << ", ";
+ } catch (std::exception& exc) {
+ LOG(FATAL) << "Constant pool can only contain NDArray and DLDataType, but got "
+ << ArgTypeCode2Str(it.type_code());
+ }
+ }
+ }
+ if (!constants.empty()) oss.seekp(-2, oss.cur);
+ oss << "]" << std::endl;
+
+ // Get the number of globals and the name of each of them.
+ oss << " Globals (#" << func_table.size() << "): [";
+ for (const auto& it : func_table) {
+ oss << it.name << ", ";
+ }
+ if (!func_map.empty()) oss.seekp(-2, oss.cur);
+ oss << "]" << std::endl;
+
+ return oss.str();
+}
+
+void Executable::SetInstructionData(Index i, Index j, ExecWord val) {
+ ICHECK_LT(i, instr_offset.size());
+ Index instr_idx = instr_offset[i];
+ ICHECK_LT(instr_idx + j, instr_data.size());
+ instr_data[instr_idx + j] = val;
+}
+
+Instruction Executable::GetInstruction(Index i) const {
+ Index offset = instr_offset[i];
+ Opcode op = static_cast<Opcode>(instr_data[offset]);
+ switch (op) {
+ case Opcode::Call: {
+ RegName dst = instr_data[offset + 1];
+ Index func_idx = instr_data[offset + 2];
+ Index num_args = instr_data[offset + 3];
+ ExecWord* args = const_cast<ExecWord*>(&instr_data[offset + 4]);
+ return Instruction::Call(func_idx, num_args, reinterpret_cast<Instruction::Arg*>(args), dst);
+ }
+ case Opcode::Ret: {
+ RegName result = instr_data[offset + 1];
+ return Instruction::Ret(result);
+ }
+ case Opcode::Goto: {
+ Index pc_offset = instr_data[offset + 1];
+ return Instruction::Goto(pc_offset);
+ }
+ case Opcode::If: {
+ RegName cond = instr_data[offset + 1];
+ Index false_offset = instr_data[offset + 2];
+ return Instruction::If(cond, false_offset);
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(op);
+ break;
+ }
+ return Instruction();
+}
+
+void SaveHeader(dmlc::Stream* strm) {
+ uint64_t header = kTVMVMBytecodeMagic;
+ strm->Write(header);
+ std::string version = TVM_VERSION;
+ strm->Write(version);
+}
+
+void LoadHeader(dmlc::Stream* strm) {
+ // Check header.
+ uint64_t header;
+ STREAM_CHECK(strm->Read(&header), "header");
+ STREAM_CHECK(header == kTVMVMBytecodeMagic, "header");
+
+ // Check version.
+ std::string version;
+ STREAM_CHECK(strm->Read(&version), "version");
+ STREAM_CHECK(version == TVM_VERSION, "version");
+}
+
+void Executable::SaveToBinary(dmlc::Stream* stream) {
+ std::string code;
+ // Initialize the stream object.
+ dmlc::MemoryStringStream strm(&code);
+
+ // Save header
+ SaveHeader(&strm);
+
+ // Global section.
+ SaveGlobalSection(&strm);
+
+ // Constant section.
+ SaveConstantSection(&strm);
+
+ // Code section.
+ SaveCodeSection(&strm);
+
+ stream->Write(code);
+}
+
- void Executable::SaveToFile(const std::string& file_name, const std::string& format) {
++void Executable::SaveToFile(const String& file_name, const String& format) {
+ std::string data;
+ dmlc::MemoryStringStream writer(&data);
+ dmlc::SeekStream* strm = &writer;
+ Executable::SaveToBinary(strm);
+ runtime::SaveBinaryToFile(file_name, data);
+}
+
+Module Executable::LoadFromBinary(void* stream) {
+ std::string code;
+ static_cast<dmlc::Stream*>(stream)->Read(&code);
+ dmlc::MemoryStringStream strm(&code);
+
+ ObjectPtr<Executable> exec = make_object<Executable>();
+
+ // Load header.
+ LoadHeader(&strm);
+
+ // Global section.
+ exec->LoadGlobalSection(&strm);
+
+ // Constant section.
+ exec->LoadConstantSection(&strm);
+
+ // Code section.
+ exec->LoadCodeSection(&strm);
+
+ return Module(exec);
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_relax.Executable")
+ .set_body_typed(Executable::LoadFromBinary);
+
- Module Executable::LoadFromFile(const std::string& file_name) {
++Module Executable::LoadFromFile(const String& file_name) {
+ std::string data;
+ runtime::LoadBinaryFromFile(file_name, &data);
+ dmlc::MemoryStringStream reader(&data);
+ dmlc::Stream* strm = &reader;
+ return Executable::LoadFromBinary(reinterpret_cast<void*>(strm));
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadfile_relax.Executable")
+ .set_body_typed(Executable::LoadFromFile);
+
+void VMFuncInfo::Save(dmlc::Stream* strm) const {
+ int32_t temp_kind = static_cast<int32_t>(kind);
+ strm->Write(temp_kind);
+ strm->Write(name);
+ strm->Write(start_instr);
+ strm->Write(end_instr);
+ strm->Write(num_args);
+ strm->Write(register_file_size);
+ strm->Write(param_names);
+}
+
+bool VMFuncInfo::Load(dmlc::Stream* strm) {
+ int32_t temp_kind;
+ if (!strm->Read(&temp_kind)) return false;
+ this->kind = static_cast<VMFuncInfo::FuncKind>(temp_kind);
+ if (!strm->Read(&name)) return false;
+ if (!strm->Read(&start_instr)) return false;
+ if (!strm->Read(&end_instr)) return false;
+ if (!strm->Read(&num_args)) return false;
+ if (!strm->Read(®ister_file_size)) return false;
+ if (!strm->Read(¶m_names)) return false;
+ return true;
+}
+
+void Executable::SaveGlobalSection(dmlc::Stream* strm) { strm->Write(func_table); }
+
+void Executable::SaveConstantSection(dmlc::Stream* strm) {
+ strm->Write(static_cast<uint64_t>(this->constants.size()));
+ for (const auto& it : this->constants) {
+ if (it.IsObjectRef<runtime::NDArray>()) {
+ strm->Write(ConstantType::kNDArray);
+ runtime::SaveDLTensor(strm, it.operator DLTensor*());
+ } else if (it.IsObjectRef<ShapeTuple>()) {
+ ShapeTuple shape = it.operator ShapeTuple();
+ strm->Write(ConstantType::kShapeTuple);
+ strm->Write(shape.size());
+ for (size_t i = 0; i < shape.size(); ++i) {
+ strm->Write(shape.at(i));
+ }
+ } else if (it.IsObjectRef<String>()) {
+ String str = it.operator String();
+ strm->Write(ConstantType::kString);
+ strm->Write(str.size());
+ for (size_t i = 0; i < str.size(); ++i) {
+ strm->Write(str.at(i));
+ }
+ } else if (it.type_code() == kDLInt) {
+ strm->Write(ConstantType::kInt);
+ strm->Write(it.value());
+ } else {
+ try {
+ strm->Write(ConstantType::kDLDataType);
+ strm->Write(it.operator DLDataType());
+ } catch (std::exception& exc) {
+ LOG(FATAL) << "Constant pool can only contain NDArray, DLDataType, and Integers but got "
+ << ArgTypeCode2Str(it.type_code());
+ }
+ }
+ }
+}
+
+void Executable::SaveCodeSection(dmlc::Stream* strm) {
+ strm->Write(instr_offset);
+ strm->Write(instr_data);
+}
+
+void Executable::LoadGlobalSection(dmlc::Stream* strm) {
+ STREAM_CHECK(strm->Read(&func_table), "Global Section");
+ // setup func map
+ for (size_t i = 0; i < func_table.size(); ++i) {
+ this->func_map[func_table[i].name] = i;
+ }
+}
+
+void Executable::LoadConstantSection(dmlc::Stream* strm) {
+ uint64_t sz;
+ // Load the number of constants.
+ STREAM_CHECK(strm->Read(&sz, sizeof(sz)), "constant");
+
+ size_t size = static_cast<size_t>(sz);
+ runtime::NDArray ndarray;
+ DLDataType dtype;
+ // Load each of the constants.
+ for (size_t i = 0; i < size; i++) {
+ int constant_type;
+ STREAM_CHECK(strm->Read(&constant_type, sizeof(constant_type)), "constant");
+ if (constant_type == ConstantType::kNDArray) {
+ ndarray.Load(strm);
+ TVMRetValue cell;
+ cell = ndarray;
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kShapeTuple) {
+ uint64_t size;
+ strm->Read(&size);
+ std::vector<ShapeTuple::index_type> data(size);
+ for (size_t i = 0; i < size; ++i) {
+ strm->Read(&(data[i]));
+ }
+ TVMRetValue cell;
+ cell = ShapeTuple(data);
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kDLDataType) {
+ strm->Read(&dtype);
+ TVMRetValue cell;
+ cell = dtype;
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kString) {
+ uint64_t size;
+ strm->Read(&size);
+ std::vector<char> data(size);
+ for (size_t i = 0; i < size; ++i) {
+ strm->Read(&(data[i]));
+ }
+ TVMRetValue cell;
+ cell = String(std::string(data.begin(), data.end()));
+ this->constants.push_back(cell);
+ } else if (constant_type == ConstantType::kInt) {
+ int64_t value;
+ strm->Read(&value);
+ TVMRetValue cell;
+ cell = value;
+ this->constants.push_back(cell);
+ } else {
+ LOG(FATAL) << "Constant pool can only contain NDArray and DLDataType, but got "
+ << ArgTypeCode2Str(constant_type) << " when loading the VM constant pool.";
+ }
+ }
+}
+
+void Executable::LoadCodeSection(dmlc::Stream* strm) {
+ STREAM_CHECK(strm->Read(&(this->instr_offset)), "instr offset");
+ STREAM_CHECK(strm->Read(&(this->instr_data)), "instr data");
+}
+
+template <typename T>
+std::string StrJoin(T* items, int offset, int cnt, std::string delim = ", ",
+ std::function<std::string(T)> repr = std::to_string) {
+ if (cnt == 0) {
+ return "";
+ }
+ std::ostringstream oss;
+ oss << repr(items[offset]);
+ for (int i = 1; i < cnt; ++i) {
+ oss << delim << repr(items[offset + i]);
+ }
+ return oss.str();
+}
+
+std::string RegNameToStr(RegName reg) {
+ if (reg == Instruction::kVoidRegister) {
+ return "%void";
+ }
+ if (reg == Instruction::kVMRegister) {
+ return "%vm";
+ }
+ return "%" + std::to_string(reg);
+}
+
+String Executable::AsText() const {
+ auto get_func_name = [&](Index index) -> std::string {
+ if (static_cast<size_t>(index) < func_table.size()) {
+ return func_table[index].name;
+ } else {
+ return "unknown_func_index(" + std::to_string(index) + ")";
+ }
+ };
+
+ auto instr_to_str = [&](Instruction::Arg arg) -> std::string {
+ // only for argument
+ switch (arg.kind()) {
+ case Instruction::ArgKind::kRegister:
+ return RegNameToStr(arg.value());
+ case Instruction::ArgKind::kImmediate:
+ return "i" + std::to_string(arg.value());
+ case Instruction::ArgKind::kConstIdx:
+ return "c[" + std::to_string(arg.value()) + "]";
+ case Instruction::ArgKind::kFuncIdx:
+ return "f[" + get_func_name(arg.value()) + "]";
+ default:
+ LOG(FATAL) << "Wrong instruction kind: " << static_cast<int>(arg.kind());
+ return "";
+ }
+ };
+
+ // print the text format
+ std::ostringstream os;
+ for (size_t fidx = 0; fidx < this->func_table.size(); ++fidx) {
+ const VMFuncInfo& gfunc = this->func_table[fidx];
+ if (gfunc.kind == VMFuncInfo::FuncKind::kPackedFunc) {
+ os << "@" << gfunc.name << " packed_func;\n\n";
+ continue;
+ }
+ if (gfunc.kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
+ os << "@" << gfunc.name << " num_inputs=" << gfunc.num_args << " vm_tir_func;\n\n";
+ continue;
+ }
+ ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
+ os << "@" << gfunc.name << ":\n";
+ size_t start_instr = gfunc.start_instr;
+ size_t end_instr = gfunc.end_instr;
+
+ for (size_t idx = start_instr; idx < end_instr; ++idx) {
+ os << " ";
+ Instruction instr = this->GetInstruction(idx);
+ switch (instr.op) {
+ case Opcode::Call: {
+ os << std::setw(6) << std::left << "call" << std::setw(16) << std::left
+ << get_func_name(instr.func_idx) << " in: " << std::setw(12) << std::left
+ << StrJoin<Instruction::Arg>(instr.args, 0, instr.num_args, ", ", instr_to_str)
+ << " dst: " << RegNameToStr(instr.dst) << "\n";
+ break;
+ }
+ case Opcode::Ret: {
+ os << std::setw(6) << std::left << "ret " << RegNameToStr(instr.result) << "\n";
+ break;
+ }
+ case Opcode::Goto: {
+ os << std::setw(6) << std::left << "goto" << instr.pc_offset << "\n";
+ break;
+ }
+ case Opcode::If: {
+ os << std::setw(6) << std::left << "If" << RegNameToStr(instr.cond) << ", "
+ << instr.false_offset << "\n";
+ break;
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(instr.op);
+ break;
+ }
+ }
+ os << "\n";
+ }
+ return String(os.str());
+}
+
+String Executable::AsPython() const {
+ auto get_func_name = [&](Index index) -> std::string {
+ if (static_cast<size_t>(index) < func_table.size()) {
+ return "\"" + func_table[index].name + "\"";
+ } else {
+ return "ib.unknown_func_index(" + std::to_string(index) + ")";
+ }
+ };
+
+ auto arg_to_py_str = [&](Instruction::Arg arg) -> std::string {
+ switch (arg.kind()) {
+ case Instruction::ArgKind::kRegister:
+ if (arg.value() == Instruction::kVMRegister) {
+ return "ib.r(vm)";
+ }
+ return "ib.r(" + std::to_string(arg.value()) + ")";
+ case Instruction::ArgKind::kImmediate:
+ return "ib.imm(" + std::to_string(arg.value()) + ")";
+ case Instruction::ArgKind::kConstIdx:
+ return "ib.c(" + std::to_string(arg.value()) + ")";
+ case Instruction::ArgKind::kFuncIdx: {
+ return "ib.f(" + get_func_name(arg.value()) + ")";
+ }
+ default:
+ LOG(FATAL) << "Wrong instruction kind: " << static_cast<int>(arg.kind());
+ return "";
+ }
+ };
+
+ // print the python format
+ std::ostringstream os;
+ os << "ib = rx.Builder()\n";
+ for (size_t fidx = 0; fidx < this->func_table.size(); ++fidx) {
+ const VMFuncInfo& gfunc = this->func_table[fidx];
+ if (gfunc.kind == VMFuncInfo::FuncKind::kPackedFunc) {
+ continue;
+ }
+ if (gfunc.kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
+ continue;
+ }
+ ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
+
+ os << "with ib.function(\"" << gfunc.name << "\", num_inputs=" << gfunc.num_args << "):\n";
+ size_t start_instr = gfunc.start_instr;
+ size_t end_instr = gfunc.end_instr;
+
+ for (size_t idx = start_instr; idx < end_instr; ++idx) {
+ Instruction instr = this->GetInstruction(idx);
+ switch (instr.op) {
+ case Opcode::Call: {
+ os << " ib.emit_call(" << get_func_name(instr.func_idx) << ", args=["
+ << StrJoin<Instruction::Arg>(instr.args, 0, instr.num_args, ", ", arg_to_py_str)
+ << "]";
+ if (instr.dst != Instruction::kVoidRegister) os << ", dst=ib.r(" << instr.dst << ")";
+ os << ")\n";
+ break;
+ }
+ case Opcode::Ret: {
+ os << " ib.emit_ret(ib.r(" << instr.result << "))\n";
+ break;
+ }
+ case Opcode::Goto: {
+ os << " ib.emit_goto(" << instr.pc_offset << ")\n";
+ break;
+ }
+ case Opcode::If: {
+ os << " ib.emit_if(ib.r(" << instr.cond << "), " << instr.false_offset << ")\n";
+ break;
+ }
+ default:
+ LOG(FATAL) << "should never hit this case: " << static_cast<int>(instr.op);
+ break;
+ }
+ }
+ }
+ return String(os.str());
+}
+
+TVM_REGISTER_GLOBAL("relax.ExecutableLoadFromFile").set_body_typed(Executable::LoadFromFile);
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --cc src/runtime/relax_vm/ndarray_cache_support.cc
index 410859d7e1,0000000000..7d1d1fc59e
mode 100644,000000..100644
--- a/src/runtime/relax_vm/ndarray_cache_support.cc
+++ b/src/runtime/relax_vm/ndarray_cache_support.cc
@@@ -1,233 -1,0 +1,233 @@@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file src/runtime/relax_vm/ndarray_cache_support.cc
+ * \brief Runtime to support ndarray cache file loading.
+ *
+ * This file provides a minimum support for ndarray cache file loading.
+ *
+ * The main focus of this implementation is to enable loading
+ * with minimum set of intermediate files while also being
+ * compatible to some of the multi-shard files that are more
+ * friendly in some of the environments.
+ *
+ * NDArray cache also provides a way to do system-wide
+ * parameter sharing across multiple VMs.
+ *
+ * There are likely other ways to load the parameters ndarray-ache.
+ * We will keep the impact minimum by puting it as a private
+ * runtime builtin provide as in this file.
+ */
+#define PICOJSON_USE_INT64
+
+#include <picojson.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/registry.h>
+
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "../../support/utils.h"
+#include "../file_utils.h"
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+/*!
+ * A NDArray cache to store pre-loaded arrays in the system.
+ */
+class NDArrayCache {
+ public:
+ static NDArrayCache* Global() {
+ static NDArrayCache* inst = new NDArrayCache();
+ return inst;
+ }
+
+ static void Update(String name, NDArray arr, bool override) {
+ NDArrayCache* pool = Global();
+ if (!override) {
+ ICHECK_EQ(pool->pool_.count(name), 0) << "Name " << name << " already exists in the cache";
+ }
+ pool->pool_.Set(name, arr);
+ }
+
+ static Optional<NDArray> Get(String name) {
+ NDArrayCache* pool = Global();
+ auto it = pool->pool_.find(name);
+ if (it != pool->pool_.end()) {
+ return (*it).second;
+ } else {
+ return NullOpt;
+ }
+ }
+
+ static void Remove(String name) {
+ NDArrayCache* pool = Global();
+ pool->pool_.erase(name);
+ }
+
+ static void Clear() { Global()->pool_.clear(); }
+
+ /*!
+ * \brief Load parameters from path and append them.
+ *
+ * \param cache_path The cache to path.
+ * \param device_type The type of device to be loaded.
+ * \param device_id The device id.
+ */
+ static void Load(const std::string& cache_path, int device_type, int device_id) {
+ DLDevice device{static_cast<DLDeviceType>(device_type), device_id};
+ std::string json_str;
+ LoadBinaryFromFile(cache_path + "/ndarray-cache.json", &json_str);
+ picojson::value json_info;
+ picojson::parse(json_info, json_str);
+ auto shard_records = json_info.get<picojson::object>()["records"].get<picojson::array>();
+
+ Map<String, NDArray> result;
+ std::string raw_data;
+ Optional<NDArray> staging_buffer;
+
+ auto fcopy_param_from_bytes = [&](NDArray param, void* data, size_t nbytes) {
+ if (device_type != kDLOpenCL) {
+ param.CopyFromBytes(data, nbytes);
+ }
+ // special handle OpenCL
+ // OpenCL runtime can create a host side memory mirror
+ // for every cl_mem that tries to copy data from host
+ // which can cause memory issue.
+ // We use a single staging buffer here
+ // that get de-allocated later
+ if (staging_buffer.defined()) {
+ size_t curr_size = runtime::GetDataSize(*(staging_buffer.value().operator->()));
+ if (curr_size < nbytes) {
+ staging_buffer = NullOpt;
+ }
+ }
+ if (!staging_buffer.defined()) {
+ staging_buffer = NDArray::Empty(param.Shape(), param->dtype, param->device);
+ }
+ NDArray staging_view = staging_buffer.value().CreateView(param.Shape(), param->dtype);
+ staging_view.CopyFromBytes(data, nbytes);
+ param.CopyFrom(staging_view);
+ TVMSynchronize(device_type, device_id, nullptr);
+ };
+
+ for (auto shard_item : shard_records) {
+ auto shard_rec = shard_item.get<picojson::object>();
+ ICHECK(shard_rec["dataPath"].is<std::string>());
+ std::string data_path = shard_rec["dataPath"].get<std::string>();
+
+ LoadBinaryFromFile(cache_path + "/" + data_path, &raw_data);
+ CHECK_EQ(shard_rec["format"].get<std::string>(), "raw-shard");
+ int64_t raw_nbytes = shard_rec["nbytes"].get<int64_t>();
+ CHECK_EQ(raw_nbytes, raw_data.length())
+ << "ValueError: Parameters are not loaded properly. Please check your parameter shards "
+ "and git lfs installation";
+
+ for (auto nd_item : shard_rec["records"].get<picojson::array>()) {
+ auto nd_rec = nd_item.get<picojson::object>();
+ CHECK(nd_rec["name"].is<std::string>());
+ String name = nd_rec["name"].get<std::string>();
+
+ std::vector<int64_t> shape;
+ for (auto value : nd_rec["shape"].get<picojson::array>()) {
+ shape.push_back(value.get<int64_t>());
+ }
+
+ DataType dtype(String2DLDataType(nd_rec["dtype"].get<std::string>()));
+ std::string encode_format = nd_rec["format"].get<std::string>();
+ int64_t offset = nd_rec["byteOffset"].get<int64_t>();
+ int64_t nbytes = nd_rec["nbytes"].get<int64_t>();
+ NDArray arr = NDArray::Empty(ShapeTuple(shape.begin(), shape.end()), dtype, device);
+
+ if (dtype == DataType::Float(32) && encode_format == "f32-to-bf16") {
+ // decode bf16 to f32
+ std::vector<uint16_t> buffer(nbytes / 2);
+ std::vector<uint32_t> decoded(nbytes / 2);
+ std::memcpy(buffer.data(), raw_data.data() + offset, nbytes);
+ for (size_t i = 0; i < buffer.size(); ++i) {
+ decoded[i] = static_cast<uint32_t>(buffer[i]) << 16;
+ }
+ fcopy_param_from_bytes(arr, decoded.data(), decoded.size() * sizeof(uint32_t));
+ } else {
+ fcopy_param_from_bytes(arr, raw_data.data() + offset, nbytes);
+ }
+ Update(name, arr, true);
+ }
+ }
+ }
+
+ private:
+ Map<String, NDArray> pool_;
+};
+
+TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.get").set_body_typed(NDArrayCache::Get);
+TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.update").set_body_typed(NDArrayCache::Update);
+TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.remove").set_body_typed(NDArrayCache::Remove);
+TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.clear").set_body_typed(NDArrayCache::Clear);
+TVM_REGISTER_GLOBAL("vm.builtin.ndarray_cache.load").set_body_typed(NDArrayCache::Load);
+
+// This param module node can be useful to get param dict in RPC mode
+// when the remote already have loaded parameters from file.
+class ParamModuleNode : public runtime::ModuleNode {
+ public:
+ const char* type_key() const final { return "param_module"; }
+
- PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final {
++ PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) final {
+ if (name == "get_params") {
+ auto params = params_;
+ return PackedFunc([params](TVMArgs args, TVMRetValue* rv) { *rv = params; });
+ } else {
+ return PackedFunc();
+ }
+ }
+
- static Array<NDArray> GetParams(const std::string& prefix, int num_params) {
++ static Array<NDArray> GetParams(const String& prefix, int num_params) {
+ Array<NDArray> params;
+ for (int i = 0; i < num_params || num_params == -1; ++i) {
+ std::string name = prefix + "_" + std::to_string(i);
+ auto opt = NDArrayCache::Get(name);
+ if (opt) {
+ params.push_back(opt.value());
+ } else {
+ if (num_params == -1) return params;
+ LOG(FATAL) << "Cannot find " << name << " in cache";
+ }
+ }
+ return params;
+ }
+
+ static Module Create(const std::string& prefix, int num_params) {
+ auto n = make_object<ParamModuleNode>();
+ n->params_ = GetParams(prefix, num_params);
+ return Module(n);
+ }
+
+ private:
+ Array<NDArray> params_;
+};
+
+TVM_REGISTER_GLOBAL("vm.builtin.param_module_from_cache").set_body_typed(ParamModuleNode::Create);
+TVM_REGISTER_GLOBAL("vm.builtin.param_array_from_cache").set_body_typed(ParamModuleNode::GetParams);
+
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --cc src/runtime/relax_vm/vm.cc
index 01497fdd7c,0000000000..9c25ca5307
mode 100644,000000..100644
--- a/src/runtime/relax_vm/vm.cc
+++ b/src/runtime/relax_vm/vm.cc
@@@ -1,1038 -1,0 +1,1038 @@@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file src/runtime/relax_vm/vm.cc
+ */
+
+#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/profiling.h>
+#include <tvm/runtime/relax_vm/vm.h>
+
+#include <optional>
+
+namespace tvm {
+namespace runtime {
+namespace relax_vm {
+
+//---------------------------------------------
+// VM Closure object
+//---------------------------------------------
+TVM_REGISTER_OBJECT_TYPE(VMClosureObj);
+
+VMClosure::VMClosure(String func_name, PackedFunc impl) {
+ auto ptr = make_object<VMClosureObj>();
+ ptr->func_name = func_name;
+ ptr->impl = std::move(impl);
+ data_ = std::move(ptr);
+}
+
+/*!
+ * \brief Create another PackedFunc with last arguments already bound to last_args.
+ * \param func The input func, can be a VMClosure or PackedFunc.
+ * \param last_args The arguments to bound to in the end of the function.
+ * \note The new function takes in arguments and append the last_args in the end.
+ */
+PackedFunc VMClosure::BindLastArgs(PackedFunc func, std::vector<TVMRetValue> last_args) {
+ return PackedFunc([func, last_args](TVMArgs args, TVMRetValue* rv) {
+ std::vector<TVMValue> values(args.size() + last_args.size());
+ std::vector<int> tcodes(args.size() + last_args.size());
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+ std::copy(args.values, args.values + args.size(), values.data());
+ std::copy(args.type_codes, args.type_codes + args.size(), tcodes.data());
+ for (size_t i = 0; i < last_args.size(); ++i) {
+ setter(i + args.size(), last_args[i]);
+ }
+ func.CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), rv);
+ });
+}
+
+//-----------------------------------------------------------
+// Utility functions.
+//-----------------------------------------------------------
+// Use the args after `starting_arg_idx` as a series of indices into `obj`,
+// indexing into nested ADTs and returning the final indexed object.
+ObjectRef IndexIntoNestedObject(ObjectRef obj, TVMArgs args, int starting_arg_idx) {
+ for (int i = starting_arg_idx; i < args.size(); i++) {
+ // the object must be an ADT to be able to index into it
+ if (!obj.as<ArrayNode>()) {
+ LOG(FATAL) << "ValueError: Attempted to index into an object that is not an ADT.";
+ }
+ int index = args[i];
+ auto arr = Downcast<Array<ObjectRef>>(obj);
+ // make sure the index is in bounds
+ if (index >= static_cast<int>(arr.size())) {
+ LOG(FATAL) << "IndexError: Invalid index (" << index << " >= " << arr.size() << ").";
+ }
+ obj = arr[index];
+ }
+ return obj;
+}
+
+NDArray ConvertNDArrayToDevice(NDArray src, const DLDevice& dev, Allocator* alloc) {
+ if (src->device.device_type == dev.device_type && src->device.device_id == dev.device_id) {
+ return src;
+ } else {
+ auto res = alloc->Empty(src.Shape(), src->dtype, dev);
+ res.CopyFrom(src);
+ return res;
+ }
+}
+
+ObjectRef ConvertObjectToDevice(ObjectRef src, const Device& dev, Allocator* alloc) {
+ if (src->IsInstance<NDArray::ContainerType>()) {
+ return ConvertNDArrayToDevice(Downcast<NDArray>(src), dev, alloc);
+ } else if (src->IsInstance<ArrayNode>()) {
+ std::vector<ObjectRef> ret;
+ auto arr = Downcast<Array<ObjectRef>>(src);
+ for (size_t i = 0; i < arr.size(); i++) {
+ ret.push_back(ConvertObjectToDevice(arr[i], dev, alloc));
+ }
+ return Array<ObjectRef>(ret.begin(), ret.end());
+ } else {
+ return src;
+ }
+}
+
+TVMRetValue ConvertArgToDevice(TVMArgValue input, Device dev, Allocator* alloc) {
+ // NOTE: NDArray::FromExternalDLTensor is not safe
+ // in terms of memory-behavior.
+ // To be extra careful, we copy DLTensor.
+ // The developer can still explicitly allocate NDArray
+ // in TVM Native API or NDArray::FromDLPack to regain zero copy behavior.
+ TVMRetValue ret;
+
+ if (input.type_code() == kTVMDLTensorHandle) {
+ DLTensor* tensor = input;
+ std::vector<int64_t> shape(tensor->shape, tensor->shape + tensor->ndim);
+ auto dst = alloc->Empty(shape, tensor->dtype, dev);
+ dst.CopyFrom(tensor);
+ ret = dst;
+ } else if (input.IsObjectRef<ObjectRef>()) {
+ ret = ConvertObjectToDevice(input.operator ObjectRef(), dev, alloc);
+ } else {
+ ret = input;
+ }
+ return ret;
+}
+
+TVMRetValue ConvertRegToDevice(TVMRetValue input, Device dev, Allocator* alloc) {
+ TVMRetValue ret;
+ if (input.IsObjectRef<ObjectRef>()) {
+ ret = ConvertObjectToDevice(input.operator ObjectRef(), dev, alloc);
+ } else {
+ ret = input;
+ }
+ return ret;
+}
+
+//-----------------------------------------------------------
+// VM implementations.
+//-----------------------------------------------------------
+/*!
+ * \brief The register type.
+ */
+using RegType = TVMRetValue;
+
+/*!
+ * \brief A representation of a stack frame.
+ *
+ * A stack frame is a record containing the information needed
+ * to restore the caller's virtual machine state after returning
+ * from a function call.
+ */
+struct VMFrame {
+ /*! \brief The return program counter. */
+ Index return_pc;
+ /*! \brief Statically allocated space for objects */
+ std::vector<RegType> register_file;
+ /*! \brief Register in caller's frame to put return value */
+ RegName caller_return_register;
+ // The following fields are used for PackedFunc call within
+ // a single function scope. The space is reused across multiple
+ // packed func calls to increase cache locality and avoid re-allocation
+ /*! \brief Temporary argument value stack for packed func call. */
+ std::vector<TVMValue> call_arg_values;
+ /*! \brief Temporary argument tcode stack for packed func call. */
+ std::vector<int> call_arg_tcodes;
+
+ VMFrame(Index pc, Index register_file_size)
+ : return_pc(pc), register_file(register_file_size), caller_return_register(0) {}
+};
+
+class VirtualMachineImpl : public VirtualMachine {
+ public:
+ //---------------------------------------------------
+ // Public facing functions overloading
+ //---------------------------------------------------
+ void LoadExecutable(ObjectPtr<Executable> exec) final;
+
+ void Init(const std::vector<Device>& devices,
+ const std::vector<AllocatorType>& alloc_types) final;
+
- PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) override;
++ PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) override;
+
+ VMClosure GetClosure(const String& func_name) final {
+ return this->GetClosureInternal(func_name, false).value();
+ }
+
+ void InvokeClosurePacked(const ObjectRef& closure_or_packedfunc, TVMArgs args,
+ TVMRetValue* rv) final;
+
+ void SetInstrument(PackedFunc instrument) final { this->instrument_ = instrument; }
+
+ //--------------------------------------------------
+ // Additional support arguments functions for VM
+ //--------------------------------------------------
+ /*!
+ * \brief Internal implementation of GetClosure which also allow none.
+ * \param func_name The name of the function.
+ * \param allow_missing Whether none is allowed.
+ * \return The result
+ */
+ Optional<VMClosure> GetClosureInternal(const String& func_name, bool allow_missing);
+
+ /*!
+ * \brief Set inputs to a function.
+ * \param func_name The function name.
+ * \param args args[offset:] are arguments to the function. If the arguments are not of the
+ * correct device for the function, they will be copied to the device.
+ * \param offset Starting offset of the arguments in \p args.
+ * \param with_param_module If set to true, the last argument will be a module and can be invoked
+ * to get the argument, this is mainly used for debugging purposes and setting composite
+ * objects. \note This interface works when using VM over RPC by internally converting NDArray in
+ * the arguments to DLTensor, which is supported in RPC where remote could only have a minimal C
+ * runtime.
+ */
+ void SetInput(std::string func_name, TVMArgs args, int offset, bool with_param_module = false);
+
+ /*!
+ * \brief Look up whether the VM has a function by the given name.
+ * \param func_name the function's name
+ * \return The function, if it exists. Logs a fatal error if not.
+ */
+ VMFuncInfo LookupVMFuncInfo(const std::string& func_name);
+
+ /*!
+ * \brief Look up whether the VM has outputs for the given function.
+ * \param func_name the function's name
+ * \return The output, if it exists. Logs a fatal error if not.
+ */
+ RegType LookupVMOutput(const std::string& func_name);
+
+ /*!
+ * \brief Fully bind the argument of a global function and save it in the env.
+ * \param func_name The global function name to be saved.
+ * \param save_name The saved name of the function.
+ * \param include_return Whether forward the return value, set it to false allows
+ * us to ignore forwarding return value, which can be helpful to do benchmarking
+ * in RPC environment when return value is complicated ADT.
+ *
+ * \param args The arguments to bound to the function.
+ * \note This function is used by RPC server to help benchmarking.
+ */
+ void SaveClosure(const String& func_name, const String& save_name, bool include_return,
+ TVMArgs args);
+ /*!
+ * \brief Internal function to invoke a closure.
+ * \param closure_or_packed The closure to be invoked.
+ * \param args The arguments to the function.
+ * \return The result value.
+ */
+ RegType InvokeClosureInternal(const ObjectRef& closure_or_packed,
+ const std::vector<RegType>& args);
+ /*!
+ * \brief Invoke a VM function by interpreting bytecode.
+ * \param fidx The function index.
+ * \param args The arguments to the function.
+ * \return The object representing the result.
+ */
+ RegType InvokeBytecode(Index fidx, const std::vector<RegType>& args);
+
+ protected:
+ /*!
+ * \brief Get function by querying all of the current module's imports.
+ * \param name The name of the function.
+ * \return The result function, can return PackedFunc(nullptr) if nothing is found.
+ */
+ PackedFunc GetFuncFromImports(const String& name) {
+ for (auto& lib : this->imports_) {
+ PackedFunc func = lib->GetFunction(name, true);
+ if (func.defined()) return func;
+ }
+ return PackedFunc(nullptr);
+ }
+ /*!
+ * \brief Initialize function pool.
+ */
+ void InitFuncPool();
+ //-------------------------------------------------
+ // Instruction interpretations.
+ //-------------------------------------------------
+ /*!
+ * \brief Push a call frame onto the call stack.
+ * \param ret_pc The program counter to return to.
+ * \param vm_func The function to be pushed to the call stack.
+ */
+ void PushFrame(Index ret_pc, const VMFuncInfo& vm_func) {
+ frames_.emplace_back(std::make_unique<VMFrame>(ret_pc, vm_func.register_file_size));
+ }
+ /*!
+ * \brief Pop a frame off the call stack.
+ */
+ void PopFrame() {
+ ICHECK_GT(frames_.size(), 0);
+ pc_ = frames_.back()->return_pc;
+ frames_.pop_back();
+ }
+ /*!
+ * \brief Write to a VM register.
+ * \param frame current vm frame.
+ * \param reg The register to write to.
+ * \param obj The object to write to.
+ */
+ void WriteRegister(VMFrame* frame, RegName reg, const RegType& obj) {
+ ICHECK_LT(reg, frame->register_file.size());
+ frame->register_file[reg] = obj;
+ }
+ /*!
+ * \brief Read a VM register.
+ * \param frame current vm frame.
+ * \param reg The register to read from.
+ * \return The value of the register.
+ */
+ RegType ReadRegister(VMFrame* frame, RegName reg) {
+ if (reg < Instruction::kBeginSpecialReg) {
+ return frame->register_file[reg];
+ }
+ RegType ret;
+ if (reg == Instruction::kVoidRegister) {
+ ret = nullptr;
+ } else {
+ ICHECK_EQ(reg, Instruction::kVMRegister);
+ // per convention, ctx ptr must be VirtualMachine* casted to void.
+ // this and VirtualMachine* may or may not be the same
+ // do first cast to VirtualMachine* then to void*
+ ret = static_cast<void*>(static_cast<VirtualMachine*>(this));
+ }
+ return ret;
+ }
+ /*!
+ * \brief Run call instruction.
+ * \param curr_frame The current frame.
+ * \param inst The call instruction.
+ */
+ virtual void RunInstrCall(VMFrame* curr_frame, Instruction inst);
+
+ /*! \brief Run VM dispatch loop. */
+ void RunLoop();
+
+ /*!
+ * \brief Retrieve the name of the function identified by the given index.
+ * \param idx The index into the VM executable function table.
+ * \return The name of the function.
+ */
+ const std::string& GetFuncName(int idx) { return exec_->func_table[idx].name; }
+
+ /*!
+ * \brief Retrieve the inputs for a function.
+ * \param func_name The name of the function.
+ * \return The function inputs.
+ */
+ const std::vector<RegType>& GetInputsFor(const std::string& func_name) {
+ return inputs_[func_name];
+ }
+
+ void ClearInputsFor(const std::string& func_name) { inputs_.erase(func_name); }
+
+ //--------------------------------------------------------
+ // Internal states for execution.
+ //--------------------------------------------------------
+ /*! \brief The loaded executable. */
+ ObjectPtr<Executable> exec_;
+ /*! \brief The global constant pool */
+ std::vector<TVMRetValue> const_pool_;
+ /*!
+ * \brief Function pool to cache functions in func_table
+ */
+ std::vector<TVMRetValue> func_pool_;
+ //--------------------------------------------------------
+ // Executor interface support
+ //--------------------------------------------------------
+ /*! \brief The function name to input register mapping. */
+ std::unordered_map<std::string, std::vector<RegType>> inputs_;
+ /*! \brief The function name to output register. */
+ std::unordered_map<std::string, RegType> outputs_;
+ /*! \brief A store of closures created by `save_function`. */
+ std::unordered_map<std::string, VMClosure> saved_closures_;
+ //------------------------------------------------------------
+ // VM Instruction execution.
+ //------------------------------------------------------------
+ /*!
+ * \brief The current stack of call frames.
+ * \note: Use unique ptr to avoid re-allocation and copy when frames_ get resized.
+ */
+ std::vector<std::unique_ptr<VMFrame>> frames_;
+ /*! \brief The virtual machine PC. */
+ Index pc_{0};
+ /*! \brief The special return register. */
+ RegType return_value_;
+ /*!\ brief instrument function. */
+ PackedFunc instrument_ = nullptr;
+};
+
+void VirtualMachineImpl::LoadExecutable(ObjectPtr<Executable> exec) {
+ this->exec_ = exec;
+ this->imports_ = exec_->imports();
+}
+
+void VirtualMachineImpl::Init(const std::vector<Device>& devices,
+ const std::vector<AllocatorType>& alloc_types) {
+ // TODO(@yuchen): support multi-device heterogeneous execution
+ ICHECK_LT(devices.size(), 3)
+ << "Currently relax vm only supports at most 2 devices (host + device)";
+ ICHECK_EQ(devices.size(), alloc_types.size());
+
+ this->devices.reserve(devices.size());
+ this->allocators.reserve(alloc_types.size());
+ for (size_t i = 0; i < devices.size(); i++) {
+ auto alloc = MemoryManager::GetOrCreateAllocator(devices[i], alloc_types[i]);
+ this->devices.push_back(devices[i]);
+ this->allocators.push_back(alloc);
+ }
+ // Setup constant sections.
+ this->const_pool_.reserve(exec_->constants.size());
+ for (const auto& constant : exec_->constants) {
+ if (constant.type_code() != kTVMNDArrayHandle) {
+ this->const_pool_.push_back(constant);
+ } else {
+ this->const_pool_.push_back(ConvertRegToDevice(constant, devices[0], allocators[0]));
+ }
+ }
+ // Setup function sections.
+ this->InitFuncPool();
+}
+
+VMFuncInfo VirtualMachineImpl::LookupVMFuncInfo(const std::string& func_name) {
+ ICHECK(exec_) << "The executable is not created yet.";
+ auto it = this->exec_->func_map.find(func_name);
+ CHECK(it != this->exec_->func_map.end()) << "ValueError: Unknown function: " << func_name;
+
+ return exec_->func_table[it->second];
+}
+
+RegType VirtualMachineImpl::LookupVMOutput(const std::string& func_name) {
+ if (!outputs_.count(func_name)) {
+ LOG(FATAL) << "ValueError: No output saved for call of \"" << func_name
+ << "\"; use `invoke_stateful` to call it first.";
+ }
+ return outputs_[func_name];
+}
+
- PackedFunc VirtualMachineImpl::GetFunction(const std::string& name,
++PackedFunc VirtualMachineImpl::GetFunction(const String& name,
+ const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "vm_initialization") {
+ // initialize the VirtualMachine, takes variable-length arguments
+ // first argument is a runtime::Module, followed by one or more device_type, device_id,
+ // and the AllocatorType associated with the device.
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ICHECK_EQ(args.size() % 3, 0);
+ std::vector<Device> devices;
+ std::vector<AllocatorType> alloc_types;
+ for (int i = 0; i < args.size(); i += 3) {
+ Device dev;
+ int device_type = args[i];
+ dev.device_type = DLDeviceType(device_type);
+ dev.device_id = args[i + 1];
+ int type = args[i + 2];
+ devices.push_back(dev);
+ alloc_types.push_back(AllocatorType(type));
+ }
+ this->Init(devices, alloc_types);
+ });
+ } else if (name == "save_function") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ ICHECK_GE(args.size(), 3);
+ this->SaveClosure(args[0], args[1], args[2],
+ TVMArgs(args.values + 3, args.type_codes + 3, args.size() - 3));
+ });
+ } else if (name == "invoke_closure") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ VMClosure clo = args[0];
+ this->InvokeClosurePacked(clo, TVMArgs(args.values + 1, args.type_codes + 1, args.size() - 1),
+ rv);
+ });
+ } else if (name == "set_instrument") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ PackedFunc func;
+ if (args[0].type_code() != kTVMPackedFuncHandle) {
+ String func_name = args[0];
+ const PackedFunc* factory = Registry::Get(func_name);
+ ICHECK(factory != nullptr) << "Cannot find factory " << func_name;
+ TVMRetValue rv;
+ factory->CallPacked(TVMArgs(args.values + 1, args.type_codes + 1, args.num_args - 1), &rv);
+ func = rv;
+ } else {
+ func = args[0];
+ }
+ this->SetInstrument(func);
+ });
+ } else if (name == "invoke_stateful") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ const auto& m = this->exec_->func_map;
+ if (m.find(func_name) == m.end()) {
+ LOG(FATAL) << "ValueError: Unknown function: " << func_name;
+ }
+ Index gf_idx = m.at(func_name);
+ if (!inputs_.count(func_name)) {
+ LOG(FATAL) << "ValueError: No inputs set for stateful call of " << func_name
+ << "; use `set_input` first.";
+ return;
+ }
+ outputs_[func_name] = this->InvokeClosureInternal(func_pool_[gf_idx], inputs_[func_name]);
+ });
+ } else if (name == "get_output_arity") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ RegType out = LookupVMOutput(func_name);
+ // use remaining args as indices
+ ObjectRef obj = IndexIntoNestedObject(out.AsObjectRef<ObjectRef>(), args, 1);
+ // after chasing through the indices, examine the final object
+ if (const auto* arr = obj.as<ArrayNode>()) {
+ *rv = static_cast<int>(arr->size());
+ } else {
+ *rv = -1;
+ }
+ });
+ } else if (name == "get_output") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ RegType out = LookupVMOutput(func_name);
+ // use remaining args as indices
+ ObjectRef obj = IndexIntoNestedObject(out.AsObjectRef<ObjectRef>(), args, 1);
+ if (obj.as<ArrayNode>()) {
+ LOG(FATAL) << "ValueError: `get_output` cannot return a tuple for RPC compatibility. "
+ "Please specify another index argument.";
+ return;
+ }
+ *rv = obj;
+ });
+ } else if (name == "set_input") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { SetInput(args[0], args, 1); });
+ } else if (name == "set_input_with_param_module") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { SetInput(args[0], args, 1, true); });
+ } else if (name == "get_function_arity") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ const VMFuncInfo& vm_func = LookupVMFuncInfo(func_name);
+ *rv = static_cast<int>(vm_func.param_names.size());
+ });
+ } else if (name == "get_function_param_name") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string func_name = args[0];
+ int index = args[1];
+ const VMFuncInfo& vm_func = LookupVMFuncInfo(func_name);
+ if (static_cast<size_t>(index) >= vm_func.param_names.size()) {
+ LOG(FATAL) << "ValueError: Invalid index for " << func_name << " (" << index << " out of "
+ << vm_func.param_names.size() << ")";
+ }
+ *rv = vm_func.param_names[index];
+ });
+ } else {
+ // default case, look up closure in VM.
+ if (Optional<VMClosure> opt = this->GetClosureInternal(name, true)) {
+ auto clo = opt.value();
+ return PackedFunc([sptr_to_self, this, clo](TVMArgs args, TVMRetValue* rv) {
+ this->InvokeClosurePacked(clo, args, rv);
+ });
+ } else {
+ return PackedFunc(nullptr);
+ }
+ }
+}
+
+void VirtualMachineImpl::SetInput(std::string func_name, TVMArgs args, int offset,
+ bool with_param_module) {
+ const auto& m = exec_->func_map;
+ if (m.find(func_name) != m.end()) {
+ Index gf_idx = m.at(func_name);
+ const VMFuncInfo& vm_func = exec_->func_table[gf_idx];
+ size_t params_num = vm_func.num_args;
+ ICHECK_EQ(args.size() - offset, params_num)
+ << "The number of provided parameters doesn't match the number of arguments for";
+ std::vector<RegType> func_args(params_num);
+
+ for (int i = offset; i < args.size(); ++i) {
+ int index = i - offset;
+ if (with_param_module && i == args.size() - 1) {
+ // call param func to get the arguments(usually corresponds to param pack.)
+ func_args[index] = (args[i].operator Module()).GetFunction("get_params")();
+ } else {
+ func_args[index] = ConvertArgToDevice(args[i], devices[0], allocators[0]);
+ }
+ }
+ inputs_[func_name] = func_args;
+ } else {
+ LOG(FATAL) << "ValueError: Unknown function: " << func_name;
+ }
+}
+
+//------------------------------------------
+// Closure handling
+//------------------------------------------
+void VirtualMachineImpl::InvokeClosurePacked(const ObjectRef& closure_or_packedfunc, TVMArgs args,
+ TVMRetValue* rv) {
+ // run packed call if it is a packed func.
+ if (auto* packed = closure_or_packedfunc.as<PackedFunc::ContainerType>()) {
+ packed->CallPacked(args, rv);
+ return;
+ }
+ // run closure call.
+ auto* clo = closure_or_packedfunc.as<VMClosureObj>();
+ ICHECK(clo != nullptr) << "Function expects a closure or PackedFunc ";
+
+ std::vector<TVMValue> values(args.size() + 1);
+ std::vector<int> tcodes(args.size() + 1);
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+ // per convention, ctx ptr must be VirtualMachine* casted to void.
+ // this and VirtualMachine* may or maynot be the same
+ // do first cast to VirtualMachine* then to void*
+ setter(0, static_cast<void*>(static_cast<VirtualMachine*>(this)));
+ std::copy(args.values, args.values + args.size(), values.begin() + 1);
+ std::copy(args.type_codes, args.type_codes + args.size(), tcodes.begin() + 1);
+ clo->impl.CallPacked(TVMArgs(values.data(), tcodes.data(), args.size() + 1), rv);
+}
+
+// internal variant version of invoke closurepacked
+RegType VirtualMachineImpl::InvokeClosureInternal(const ObjectRef& closure_or_packed,
+ const std::vector<RegType>& args) {
+ RegType ret;
+ auto* packed = closure_or_packed.as<PackedFunc::ContainerType>();
+ auto* clo = closure_or_packed.as<VMClosureObj>();
+ int clo_offset = clo != nullptr ? 1 : 0;
+ std::vector<TVMValue> values(args.size() + clo_offset);
+ std::vector<int> tcodes(args.size() + clo_offset);
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+
+ if (clo != nullptr) {
+ setter(0, static_cast<void*>(static_cast<VirtualMachine*>(this)));
+ }
+ for (size_t i = 0; i < args.size(); ++i) {
+ setter(i + clo_offset, args[i]);
+ }
+
+ if (packed != nullptr) {
+ packed->CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), &ret);
+ } else {
+ ICHECK(clo != nullptr);
+ clo->impl.CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), &ret);
+ }
+ return ret;
+}
+
+void VirtualMachineImpl::SaveClosure(const String& func_name, const String& save_name,
+ bool include_return, TVMArgs args) {
+ VMClosure clo = this->GetClosure(func_name);
+ std::vector<RegType> inputs(args.size());
+ for (int i = 0; i < args.size(); ++i) {
+ inputs[i] = ConvertArgToDevice(args[i], this->devices[0], this->allocators[0]);
+ }
+ PackedFunc impl = VMClosure::BindLastArgs(clo->impl, inputs);
+ if (!include_return) {
+ impl = PackedFunc([impl](TVMArgs args, TVMRetValue* rv) {
+ TVMRetValue temp;
+ impl.CallPacked(args, &temp);
+ });
+ }
+ saved_closures_[save_name] = VMClosure(save_name, impl);
+}
+
+Optional<VMClosure> VirtualMachineImpl::GetClosureInternal(const String& func_name,
+ bool allow_missing) {
+ // look up saved closures.
+ auto saved_it = saved_closures_.find(func_name);
+ if (saved_it != saved_closures_.end()) {
+ return saved_it->second;
+ }
+ auto it = exec_->func_map.find(func_name);
+ if (it == exec_->func_map.end()) {
+ if (allow_missing) return NullOpt;
+ LOG(FATAL) << "ValueError: Unknown function: " << func_name;
+ }
+
+ Index gf_idx = it->second;
+ const VMFuncInfo& finfo = exec_->func_table[gf_idx];
+
+ if (finfo.kind == VMFuncInfo::FuncKind::kVMFunc) {
+ // NOTE: should not capture strong ref to self and avoid cyclic ref.
+ auto impl = PackedFunc([gf_idx](TVMArgs args, TVMRetValue* rv) {
+ // Per convention, ctx ptr is a VirtualMachine*
+ VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].operator void*());
+
+ std::vector<RegType> inputs(args.size() - 1);
+ for (size_t i = 0; i < inputs.size(); ++i) {
+ inputs[i] = args[i + 1];
+ }
+ *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx, inputs);
+ });
+ return VMClosure(func_name, impl);
+ } else {
+ ICHECK(finfo.kind == VMFuncInfo::FuncKind::kVMTIRFunc)
+ << "Cannot support closure with function kind " << static_cast<int>(finfo.kind);
+ PackedFunc tir_func = GetFuncFromImports("__vmtir__" + finfo.name);
+ ICHECK(tir_func != nullptr) << "Cannot find underlying compiled tir function of VMTIRFunc "
+ << finfo.name;
+ auto impl = PackedFunc([this, finfo, tir_func](TVMArgs args, TVMRetValue* rv) {
+ // Per convention, ctx ptr is a VirtualMachine*
+ VirtualMachine* ctx_ptr = static_cast<VirtualMachine*>(args[0].operator void*());
+ ICHECK(ctx_ptr == this);
+ ICHECK_EQ(args.size() - 1, finfo.num_args)
+ << "Function " << finfo.name << " expects " << finfo.num_args << " arguments";
+ ICHECK_GE(finfo.register_file_size, finfo.num_args + 1);
+ std::vector<TVMRetValue> reg_file(finfo.register_file_size);
+ for (int64_t i = 0; i < finfo.num_args; ++i) {
+ reg_file[i] = args[i + 1];
+ }
+ void* reg_anylist_handle = reg_file.data();
+ void* const_anylist_handle = this->const_pool_.data();
+ void* func_anylist_handle = this->func_pool_.data();
+ tir_func(static_cast<void*>(ctx_ptr), reg_anylist_handle, const_anylist_handle,
+ func_anylist_handle);
+ // Return value always stored after inputs.
+ *rv = reg_file[finfo.num_args];
+ });
+ return VMClosure(func_name, impl);
+ }
+}
+
+//--------------------------------------------------------------------
+// Instruction interpretations.
+//--------------------------------------------------------------------
+RegType VirtualMachineImpl::InvokeBytecode(Index gf_idx, const std::vector<RegType>& args) {
+ const VMFuncInfo& gfunc = exec_->func_table[gf_idx];
+ ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
+
+ // Get the curr instr which might be a potential caller.
+ Instruction curr_instr = exec_->GetInstruction(pc_);
+ PushFrame(this->pc_, gfunc);
+ // Get new frame and set the caller info.
+ VMFrame* curr_frame = frames_.back().get();
+ if (curr_instr.op == Opcode::Call) {
+ curr_frame->caller_return_register = curr_instr.dst;
+ }
+
+ // load arguments to the register file
+ ICHECK_EQ(static_cast<size_t>(gfunc.num_args), args.size())
+ << "ValueError: Invoking function " << gfunc.name << " requires " << gfunc.num_args
+ << " inputs but only " << args.size() << " inputs are provided.";
+ for (size_t i = 0; i < args.size(); ++i) {
+ WriteRegister(frames_.back().get(), i, args[i]);
+ }
+ // set program counter
+ pc_ = gfunc.start_instr;
+ RunLoop();
+ return return_value_;
+}
+
+void VirtualMachineImpl::InitFuncPool() {
+ func_pool_.resize(exec_->func_table.size());
+
+ for (size_t func_index = 0; func_index < exec_->func_table.size(); ++func_index) {
+ const VMFuncInfo& info = exec_->func_table[func_index];
+ if (info.kind == VMFuncInfo::FuncKind::kPackedFunc) {
+ // only look through imports first
+ PackedFunc func = GetFuncFromImports(info.name);
+ if (!func.defined()) {
+ const PackedFunc* p_func = Registry::Get(info.name);
+ if (p_func != nullptr) func = *(p_func);
+ }
+ ICHECK(func.defined())
+ << "Error: Cannot find PackedFunc " << info.name
+ << " in either Relax VM kernel library, or in TVM runtime PackedFunc registry, or in "
+ "global Relax functions of the VM executable";
+ func_pool_[func_index] = func;
+
+ } else {
+ ICHECK(info.kind == VMFuncInfo::FuncKind::kVMFunc ||
+ info.kind == VMFuncInfo::FuncKind::kVMTIRFunc);
+ auto clo = this->GetClosure(info.name);
+ func_pool_[func_index] = clo;
+ }
+ }
+}
+
+void VirtualMachineImpl::RunInstrCall(VMFrame* curr_frame, Instruction instr) {
+ DLOG(INFO) << "\n pc = " << pc_ << ", execute: " << GetFuncName(instr.func_idx);
+ int args_begin_offset = instrument_ != nullptr ? 4 : 0;
+ // Use the call arg stack from the current frame to increase reuse
+ // and avoid re-allocation
+ curr_frame->call_arg_values.resize(args_begin_offset + instr.num_args);
+ curr_frame->call_arg_tcodes.resize(args_begin_offset + instr.num_args);
+
+ // NOTE: no changes and resize to those vector ref(otherwise can leads to segfault)
+ // in the remainder part of the function.
+ std::vector<TVMValue>& values = curr_frame->call_arg_values;
+ std::vector<int>& tcodes = curr_frame->call_arg_tcodes;
+
+ runtime::TVMArgsSetter setter(values.data(), tcodes.data());
+ for (Index i = 0; i < instr.num_args; ++i) {
+ Instruction::Arg arg = instr.args[i];
+ int arg_index = args_begin_offset + i;
+ switch (arg.kind()) {
+ case Instruction::ArgKind::kRegister: {
+ setter(arg_index, ReadRegister(curr_frame, arg.value()));
+ break;
+ }
+ case Instruction::ArgKind::kImmediate: {
+ setter(arg_index, arg.value());
+ break;
+ }
+ case Instruction::ArgKind::kConstIdx: {
+ setter(arg_index, this->const_pool_[arg.value()]);
+ break;
+ }
+ case Instruction::ArgKind::kFuncIdx: {
+ ICHECK_LT(static_cast<size_t>(arg.value()), this->func_pool_.size());
+ setter(arg_index, this->func_pool_[arg.value()]);
+ break;
+ }
+ default: {
+ LOG(FATAL) << "ValueError: Unknown argument kind: " << int(arg.kind());
+ }
+ }
+ }
+ TVMArgs args(values.data() + args_begin_offset, tcodes.data() + args_begin_offset,
+ instr.num_args);
+ TVMRetValue ret;
+
+ ICHECK_LT(static_cast<size_t>(instr.func_idx), this->func_pool_.size());
+
+ if (instrument_ == nullptr) {
+ this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);
+ } else {
+ // insert light-weight instrument callback
+ setter(0, func_pool_[instr.func_idx]);
+ setter(1, GetFuncName(instr.func_idx));
+ setter(2, true);
+ setter(3, nullptr);
+ TVMRetValue rv;
+ // store dtype to str since py callback cannot handle dtype atm.
+ std::vector<std::unique_ptr<std::string>> temp_dtype;
+ for (int i = 0; i < instr.num_args; ++i) {
+ if (tcodes[i + args_begin_offset] == kTVMDataType) {
+ std::string str_dtype = args[i];
+ temp_dtype.emplace_back(std::make_unique<std::string>(str_dtype));
+ setter(i + args_begin_offset, *temp_dtype.back());
+ }
+ }
+ int ret_kind = static_cast<int>(VMInstrumentReturnKind::kNoOp);
+ instrument_.CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), &rv);
+ if (rv.type_code() == kDLInt) {
+ ret_kind = rv;
+ }
+ if (ret_kind != static_cast<int>(VMInstrumentReturnKind::kSkipRun)) {
+ this->InvokeClosurePacked(func_pool_[instr.func_idx], args, &ret);
+ setter(2, false);
+ setter(3, ret);
+ instrument_.CallPacked(TVMArgs(values.data(), tcodes.data(), values.size()), &rv);
+ }
+ }
+
+ // save the return value to the register
+ // saving to special register is a NOP
+ if (instr.dst < Instruction::kBeginSpecialReg) {
+ WriteRegister(curr_frame, instr.dst, ret);
+ }
+ // increment pc
+ pc_++;
+}
+
+void VirtualMachineImpl::RunLoop() {
+ VMFrame* curr_frame = frames_.back().get();
+
+ while (true) {
+ ICHECK_LT(static_cast<size_t>(pc_), exec_->instr_offset.size()) << "run into invalid section";
+ Instruction instr = exec_->GetInstruction(pc_);
+ switch (instr.op) {
+ case Opcode::Call: {
+ this->RunInstrCall(curr_frame, instr);
+ break;
+ }
+ case Opcode::Ret: {
+ // If we have hit the point from which we started
+ // running, we should return to the caller breaking
+ // the dispatch loop.
+ return_value_ = ReadRegister(curr_frame, instr.result);
+ RegName caller_return_register = curr_frame->caller_return_register;
+ PopFrame();
+ if (frames_.size() == 0) {
+ // directly return if no frame in the call stack.
+ } else {
+ // return from a local call.
+ // Update the current frame to be the parent frame.
+ curr_frame = frames_.back().get();
+ WriteRegister(curr_frame, caller_return_register, return_value_);
+ }
+ return;
+ }
+ case Opcode::Goto: {
+ pc_ += instr.pc_offset;
+ break;
+ }
+ case Opcode::If: {
+ int64_t cond_val = ReadRegister(curr_frame, instr.cond);
+ if (cond_val != 0) {
+ pc_++;
+ } else {
+ ICHECK_GT(instr.false_offset, 1);
+ pc_ += instr.false_offset;
+ }
+ break;
+ }
+ }
+ }
+}
+
+ObjectPtr<VirtualMachine> VirtualMachine::Create() { return make_object<VirtualMachineImpl>(); }
+
+//----------------------------------------------------------------
+// Profiler can be optionally disabled via a macro to reduce dep.
+//----------------------------------------------------------------
+#if TVM_RELAX_VM_ENABLE_PROFILER
+
+/*!
+ * \brief An extension of VirtualMachineImpl to support per-op profiling
+ * It overrides RunInstrCall to add instrumentations around it.
+ */
+class VirtualMachineProfiler : public VirtualMachineImpl {
+ public:
- PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) override {
++ PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) override {
+ if (name == "profile") {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ std::string f_name = args[0];
+ VMClosure clo = this->GetClosure(f_name);
+
+ std::vector<Device> devices;
+ for (auto dev : this->devices) {
+ if (dev.device_type > 0) {
+ devices.push_back(dev);
+ }
+ }
+
+ prof_ = profiling::Profiler(devices, {}, {{String("Executor"), String("VM")}});
+
+ auto inputs = GetInputsFor(f_name);
+
+ bool clear_inputs = false;
+ if (inputs.size() == 0) {
+ ICHECK(args.num_args > 1) << "No input is provided";
+ TVMArgs f_args(args.values + 1, args.type_codes + 1, args.num_args - 1);
+ SetInput(f_name, args, 1);
+ inputs = GetInputsFor(f_name);
+ clear_inputs = true;
+ } else {
+ ICHECK_EQ(args.num_args, 1) << "Inputs are already provided by set_input.";
+ }
+
+ // warmup
+ this->InvokeClosureInternal(clo, inputs);
+
+ prof_->Start();
+ this->InvokeClosureInternal(clo, inputs);
+ prof_->Stop();
+
+ // Return the report as json, since profiling::Report object is not supported by RPC
+ std::string report_json = prof_->Report()->AsJSON();
+ *rv = report_json;
+
+ prof_ = std::nullopt; // releases hardware counters
+ if (clear_inputs) {
+ // SetInput modifies the internal states of VM. Undo the change after profiling.
+ ClearInputsFor(f_name);
+ }
+ });
+ } else {
+ return VirtualMachineImpl::GetFunction(name, sptr_to_self);
+ }
+ }
+
+ protected:
+ void RunInstrCall(VMFrame* curr_frame, Instruction inst) override {
+ bool profiling = false;
+ if (prof_ && prof_->IsRunning()) {
+ auto f_name = GetFuncName(inst.func_idx);
+ std::optional<Device> dev;
+ std::vector<NDArray> arrs;
+
+ auto f_check_ndarray_arg = [&dev, &arrs](const RegType& arg) {
+ if (arg.type_code() == kTVMNDArrayHandle) {
+ NDArray arr = arg;
+ dev = arr->device;
+ arrs.push_back(arr);
+ }
+ };
+
+ for (Index i = 0; i < inst.num_args; ++i) {
+ Instruction::Arg arg = inst.args[i];
+ if (arg.kind() == Instruction::ArgKind::kRegister) {
+ auto reg = ReadRegister(curr_frame, arg.value());
+ f_check_ndarray_arg(reg);
+ } else if (arg.kind() == Instruction::ArgKind::kConstIdx) {
+ const auto& const_val = this->const_pool_[arg.value()];
+ f_check_ndarray_arg(const_val);
+ }
+ }
+
+ std::unordered_map<std::string, ObjectRef> metrics;
+ metrics["Argument Shapes"] = profiling::ShapeString(arrs);
+
+ // If a suitable device is found, enable profiling.
+ if (dev) {
+ profiling = true;
+ prof_->StartCall(f_name, *dev, metrics);
+ }
+ }
+
+ VirtualMachineImpl::RunInstrCall(curr_frame, inst);
+
+ if (profiling) {
+ prof_->StopCall();
+ }
+ }
+
+ private:
+ std::optional<profiling::Profiler> prof_;
+};
+
+ObjectPtr<VirtualMachine> VirtualMachine::CreateProfiler() {
+ return make_object<VirtualMachineProfiler>();
+}
+
+#else
+ObjectPtr<VirtualMachine> VirtualMachine::CreateProfiler() {
+ LOG(FATAL) << "Profiler support is disabled";
+ return nullptr;
+}
+#endif // TVM_RELAX_VM_ENABLE_PROFILER
+} // namespace relax_vm
+} // namespace runtime
+} // namespace tvm
diff --cc src/target/source/codegen_webgpu.cc
index 538acb0a35,4d1d834c7f..1702699ac2
--- a/src/target/source/codegen_webgpu.cc
+++ b/src/target/source/codegen_webgpu.cc
@@@ -663,18 -502,12 +659,18 @@@ class WebGPUSourceModuleNode final : pu
stream->Write(smap_);
}
- std::string GetSource(const std::string& format) final {
+ String GetSource(const String& format) final {
- std::ostringstream os;
- for (auto kv : smap_) {
- os << kv.second;
+ if (format == "func_info") {
+ std::ostringstream stream;
+ dmlc::JSONWriter(&stream).Write(fmap_);
+ return stream.str();
+ } else {
+ std::ostringstream os;
+ for (auto kv : smap_) {
+ os << kv.second;
+ }
+ return os.str();
}
- return os.str();
}
private:
diff --cc web/emcc/webgpu_runtime.cc
index 6c27a8207e,a4facc8f77..d0ec39efa1
--- a/web/emcc/webgpu_runtime.cc
+++ b/web/emcc/webgpu_runtime.cc
@@@ -159,35 -159,7 +159,35 @@@ class WebGPUModuleNode final : public r
const char* type_key() const final { return "webgpu"; }
- PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final {
+ PackedFunc GetFunction(const String& name, const ObjectPtr<Object>& sptr_to_self) final {
+ // special function
+ if (name == "webgpu.get_fmap") {
+ return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
+ std::ostringstream os;
+ dmlc::JSONWriter writer(&os);
+ writer.Write(fmap_);
+ *rv = os.str();
+ });
+ } else if (name == "webgpu.get_shader") {
+ return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
+ std::string name = args[0];
+ auto it = smap_.find(name);
+ ICHECK(it != smap_.end()) << "Cannot find code " << name;
+ *rv = it->second;
+ });
+ } else if (name == "webgpu.update_prebuild") {
+ return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
+ std::string name = args[0];
+ PackedFunc func = args[1];
+ prebuild_[name] = func;
+ });
+ }
+ // check prebuild cache
+ auto prebuild_it = prebuild_.find(name);
+ if (prebuild_it != prebuild_.end()) {
+ return prebuild_it->second;
+ }
+
auto it = smap_.find(name);
if (it != smap_.end()) {
FunctionInfo info = fmap_.at(name);
@@@ -201,9 -173,7 +201,9 @@@
}
}
- void SaveToFile(const String& file_name, const String& format) final {
+ int GetPropertyMask() const final { return ModulePropertyMask::kBinarySerializable; };
+
- void SaveToFile(const std::string& file_name, const std::string& format) final {
++ void SaveToFile(const String& file_name, const std::string& format) final {
LOG(FATAL) << "Not implemented";
}