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 2025/02/16 00:29:13 UTC
(tvm) branch refactor updated (878b6599b7 -> 6ce1fdb225)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch refactor
in repository https://gitbox.apache.org/repos/asf/tvm.git
omit 878b6599b7 Fix gpu ci jenkins
omit c63f29846a Update jenkins to only rerun in last build
omit 33ea57bb80 Remove cpptest on gpu stage as it can be very slow to rebuild libtvm, cpptest should only run on cpu and test most basic stuffs Use python test for broader coverage.
omit cd135dc348 remove astext
omit 538d2c01a7 cleanup
omit 4e4ab3c471 Cleanup relay related python components
omit e313e46789 remove topi schedule and keep compute
omit 13f300e710 Cleanup the runtime tests
omit e3f69f412f Fix the system lib in build tir
omit 247252bb4c [REFACTOR] Phase out relay python components
add 21095b98bb [Relax][PyTorch] Add support for abs, ceil, erf, floor, log ops and refactor unary tests (#17622)
add ccaa534b2c [REFACTOR] Phase out relay python components (#17656)
new 6ce1fdb225 [REFACTOR] Phase out te.schedule python components
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (878b6599b7)
\
N -- N -- N refs/heads/refactor (6ce1fdb225)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
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:
apps/android_rpc/tests/android_rpc_test.py | 61 +-
apps/ios_rpc/tests/ios_rpc_test.py | 35 +-
docs/how_to/tutorials/cross_compilation_and_rpc.py | 18 +-
golang/sample/deploy.py | 56 -
jvm/README.md | 29 -
jvm/core/src/test/scripts/test_add_cpu.py | 43 -
jvm/core/src/test/scripts/test_add_gpu.py | 58 -
python/tvm/contrib/peak.py | 394 -------
python/tvm/contrib/sparse.py | 204 ----
python/tvm/contrib/tedd.py | 798 -------------
python/tvm/driver/build_module.py | 246 +---
python/tvm/exec/measure_peak.py | 52 -
python/tvm/relax/frontend/torch/fx_translator.py | 5 +
python/tvm/relax/vm_build.py | 20 +-
python/tvm/te/__init__.py | 12 +-
python/tvm/te/autodiff.py | 67 --
python/tvm/te/hybrid/__init__.py | 101 --
python/tvm/te/hybrid/calls.py | 183 ---
python/tvm/te/hybrid/module.py | 113 --
python/tvm/te/hybrid/parser.py | 658 -----------
python/tvm/te/hybrid/preprocessor.py | 120 --
python/tvm/te/hybrid/runtime.py | 175 ---
python/tvm/te/hybrid/utils.py | 103 --
python/tvm/te/schedule.py | 665 -----------
python/tvm/te/tensor.py | 10 -
python/tvm/te/tensor_intrin.py | 146 ---
python/tvm/testing/utils.py | 3 +-
python/tvm/tir/buffer.py | 23 -
python/tvm/topi/__init__.py | 4 -
python/tvm/topi/argwhere.py | 197 ----
python/tvm/topi/nn/conv2d.py | 85 --
python/tvm/topi/random/__init__.py | 22 -
python/tvm/topi/random/kernel.py | 657 -----------
python/tvm/topi/sparse_fill_empty_rows.py | 109 --
python/tvm/topi/transform.py | 30 -
python/tvm/topi/unique.py | 234 ----
python/tvm/topi/vision/__init__.py | 25 -
python/tvm/topi/vision/nms.py | 1183 --------------------
python/tvm/topi/vision/nms_util.py | 338 ------
python/tvm/topi/vision/rcnn/__init__.py | 22 -
python/tvm/topi/vision/rcnn/proposal.py | 448 --------
python/tvm/topi/vision/rcnn/roi_align.py | 228 ----
python/tvm/topi/vision/rcnn/roi_pool.py | 95 --
python/tvm/topi/vision/reorg.py | 42 -
python/tvm/topi/vision/ssd/__init__.py | 22 -
python/tvm/topi/vision/ssd/multibox.py | 369 ------
python/tvm/utils/__init__.py | 19 -
python/tvm/utils/roofline/__init__.py | 279 -----
python/tvm/utils/roofline/cuda.py | 407 -------
python/tvm/utils/roofline/registry.py | 111 --
python/tvm/utils/roofline/x86.py | 331 ------
.../test_minimal_target_codegen_llvm.py | 22 +-
.../test_runtime_ndarray.py | 22 -
.../python/codegen/test_target_codegen_aarch64.py | 59 +-
tests/python/codegen/test_target_codegen_arm.py | 19 +-
tests/python/codegen/test_target_codegen_bool.py | 25 +-
tests/python/codegen/test_target_codegen_c_host.py | 118 +-
.../codegen/test_target_codegen_cross_llvm.py | 14 +-
tests/python/codegen/test_target_codegen_cuda.py | 539 +++------
tests/python/codegen/test_target_codegen_device.py | 49 +-
tests/python/codegen/test_target_codegen_extern.py | 25 +-
.../python/codegen/test_target_codegen_hexagon.py | 58 +-
tests/python/codegen/test_target_codegen_llvm.py | 433 +++----
tests/python/codegen/test_target_codegen_opencl.py | 9 +-
tests/python/codegen/test_target_codegen_rocm.py | 75 +-
tests/python/codegen/test_target_codegen_vulkan.py | 247 +---
tests/python/codegen/test_target_codegen_x86.py | 6 +-
tests/python/contrib/test_cblas.py | 16 +-
tests/python/contrib/test_dlpack.py | 3 +-
tests/python/contrib/test_gemm_acc16.py | 105 --
tests/python/contrib/test_gemm_acc32_vnni.py | 115 --
.../test_hexagon/conv2d/test_conv2d_blocked.py | 207 ----
.../test_hexagon/conv2d/test_conv2d_conv2d.py | 252 -----
.../test_hexagon/test_2d_physical_buffers.py | 41 +-
tests/python/contrib/test_hexagon/test_launcher.py | 12 +-
.../contrib/test_hexagon/test_maxpool2d_blocked.py | 158 ---
tests/python/contrib/test_hipblas.py | 6 +-
tests/python/contrib/test_miopen.py | 136 ---
tests/python/contrib/test_mps.py | 26 +-
tests/python/contrib/test_random.py | 9 +-
tests/python/contrib/test_rocblas.py | 6 +-
tests/python/contrib/test_sort.py | 41 +-
tests/python/contrib/test_sparse.py | 123 --
tests/python/relax/test_frontend_from_fx.py | 795 ++++---------
tests/python/runtime/test_runtime_dlpack.py | 4 +-
tests/python/runtime/test_runtime_measure.py | 3 +-
tests/python/runtime/test_runtime_module_export.py | 213 +---
tests/python/runtime/test_runtime_module_load.py | 16 +-
.../python/runtime/test_runtime_module_property.py | 12 +-
tests/python/runtime/test_runtime_rpc.py | 3 +-
tests/python/runtime/test_runtime_trace.py | 21 +-
tests/python/te/test_te_autodiff.py | 351 ------
tests/python/te/test_te_build_lower.py | 65 --
tests/python/te/test_te_group.py | 90 --
tests/python/te/test_te_hybrid_script.py | 872 ---------------
tests/python/te/test_te_schedule.py | 382 -------
.../python/te/test_te_schedule_bound_inference.py | 512 ---------
.../te/test_te_schedule_bound_inference_tiling.py | 62 -
tests/python/te/test_te_schedule_graph.py | 142 ---
tests/python/te/test_te_schedule_lstm.py | 91 --
tests/python/te/test_te_schedule_ops.py | 695 ------------
...te_schedule_postproc_rewrite_for_tensor_core.py | 231 ----
tests/python/te/test_te_schedule_tensor_core.py | 461 --------
tests/python/te/test_te_schedule_tensorize.py | 392 -------
tests/python/te/test_te_tensor.py | 202 +---
tests/python/te/test_te_transform_layout.py | 592 ----------
.../tir-analysis/test_tir_analysis_usedef.py | 36 -
.../test_tir_analysis_verify_gpu_code.py | 434 -------
.../test_tir_analysis_verify_memory.py | 121 --
tests/python/tir-base/test_lower_build.py | 17 -
tests/python/tir-base/test_tir_buffer.py | 79 --
tests/python/tir-base/test_tir_intrin.py | 68 +-
tests/python/tir-base/test_tir_ir_builder.py | 565 ----------
.../test_tir_transform_compact_buffer_region.py | 14 -
.../test_tir_transform_convert_blocks_to_opaque.py | 9 -
.../test_tir_transform_flatten_buffer.py | 13 -
.../tir-transform/test_tir_transform_hoist_if.py | 59 -
.../test_tir_transform_inject_copy_intrin.py | 124 --
.../test_tir_transform_inject_rolling_buffer.py | 177 ---
...test_tir_transform_instrument_bound_checkers.py | 608 ----------
.../test_tir_transform_loop_partition.py | 325 ------
...t_tir_transform_lower_cross_thread_reduction.py | 15 -
.../test_tir_transform_lower_init_block.py | 9 -
.../test_tir_transform_lower_intrin.py | 4 +-
.../test_tir_transform_lower_opaque_block.py | 9 -
.../test_tir_transform_lower_warp_memory.py | 356 ------
.../test_tir_transform_make_packed_api.py | 26 -
...form_merge_dynamic_shared_memory_allocations.py | 304 -----
...sform_merge_static_shared_memory_allocations.py | 203 ----
.../test_tir_transform_narrow_datatype.py | 41 -
...sform_plan_update_buffer_allocation_location.py | 11 -
.../tir-transform/test_tir_transform_simplify.py | 64 --
.../test_tir_transform_split_host_device.py | 39 -
.../test_tir_transform_storage_flatten.py | 63 --
.../test_tir_transform_storage_rewrite.py | 320 ------
.../test_tir_transform_thread_sync.py | 61 -
.../test_tir_transform_unify_thread_binding.py | 11 -
.../test_tir_transform_unroll_loop.py | 18 -
.../tir-transform/test_tir_transform_vectorize.py | 25 +-
web/tests/python/webgpu_rpc_test.py | 3 +-
web/tests/python/websock_rpc_test.py | 91 --
141 files changed, 1012 insertions(+), 22148 deletions(-)
delete mode 100644 golang/sample/deploy.py
delete mode 100644 jvm/core/src/test/scripts/test_add_cpu.py
delete mode 100644 jvm/core/src/test/scripts/test_add_gpu.py
delete mode 100644 python/tvm/contrib/peak.py
delete mode 100644 python/tvm/contrib/sparse.py
delete mode 100644 python/tvm/contrib/tedd.py
delete mode 100644 python/tvm/exec/measure_peak.py
delete mode 100644 python/tvm/te/autodiff.py
delete mode 100644 python/tvm/te/hybrid/__init__.py
delete mode 100644 python/tvm/te/hybrid/calls.py
delete mode 100644 python/tvm/te/hybrid/module.py
delete mode 100644 python/tvm/te/hybrid/parser.py
delete mode 100644 python/tvm/te/hybrid/preprocessor.py
delete mode 100644 python/tvm/te/hybrid/runtime.py
delete mode 100644 python/tvm/te/hybrid/utils.py
delete mode 100644 python/tvm/te/schedule.py
delete mode 100644 python/tvm/te/tensor_intrin.py
delete mode 100644 python/tvm/topi/argwhere.py
delete mode 100644 python/tvm/topi/random/__init__.py
delete mode 100644 python/tvm/topi/random/kernel.py
delete mode 100644 python/tvm/topi/sparse_fill_empty_rows.py
delete mode 100644 python/tvm/topi/vision/__init__.py
delete mode 100644 python/tvm/topi/vision/nms.py
delete mode 100644 python/tvm/topi/vision/nms_util.py
delete mode 100644 python/tvm/topi/vision/rcnn/__init__.py
delete mode 100644 python/tvm/topi/vision/rcnn/proposal.py
delete mode 100644 python/tvm/topi/vision/rcnn/roi_align.py
delete mode 100644 python/tvm/topi/vision/rcnn/roi_pool.py
delete mode 100644 python/tvm/topi/vision/reorg.py
delete mode 100644 python/tvm/topi/vision/ssd/__init__.py
delete mode 100644 python/tvm/topi/vision/ssd/multibox.py
delete mode 100644 python/tvm/utils/__init__.py
delete mode 100644 python/tvm/utils/roofline/__init__.py
delete mode 100644 python/tvm/utils/roofline/cuda.py
delete mode 100644 python/tvm/utils/roofline/registry.py
delete mode 100644 python/tvm/utils/roofline/x86.py
delete mode 100644 tests/python/contrib/test_gemm_acc16.py
delete mode 100644 tests/python/contrib/test_gemm_acc32_vnni.py
delete mode 100644 tests/python/contrib/test_hexagon/conv2d/test_conv2d_blocked.py
delete mode 100644 tests/python/contrib/test_hexagon/conv2d/test_conv2d_conv2d.py
delete mode 100644 tests/python/contrib/test_hexagon/test_maxpool2d_blocked.py
delete mode 100644 tests/python/contrib/test_miopen.py
delete mode 100644 tests/python/contrib/test_sparse.py
delete mode 100644 tests/python/te/test_te_autodiff.py
delete mode 100644 tests/python/te/test_te_build_lower.py
delete mode 100644 tests/python/te/test_te_group.py
delete mode 100644 tests/python/te/test_te_hybrid_script.py
delete mode 100644 tests/python/te/test_te_schedule.py
delete mode 100644 tests/python/te/test_te_schedule_bound_inference.py
delete mode 100644 tests/python/te/test_te_schedule_bound_inference_tiling.py
delete mode 100644 tests/python/te/test_te_schedule_graph.py
delete mode 100644 tests/python/te/test_te_schedule_lstm.py
delete mode 100644 tests/python/te/test_te_schedule_ops.py
delete mode 100644 tests/python/te/test_te_schedule_postproc_rewrite_for_tensor_core.py
delete mode 100644 tests/python/te/test_te_schedule_tensor_core.py
delete mode 100644 tests/python/te/test_te_schedule_tensorize.py
delete mode 100644 tests/python/te/test_te_transform_layout.py
delete mode 100644 tests/python/tir-analysis/test_tir_analysis_usedef.py
delete mode 100644 tests/python/tir-analysis/test_tir_analysis_verify_gpu_code.py
delete mode 100644 tests/python/tir-analysis/test_tir_analysis_verify_memory.py
delete mode 100644 tests/python/tir-base/test_tir_ir_builder.py
delete mode 100644 tests/python/tir-transform/test_tir_transform_inject_copy_intrin.py
delete mode 100644 tests/python/tir-transform/test_tir_transform_instrument_bound_checkers.py
delete mode 100644 tests/python/tir-transform/test_tir_transform_lower_warp_memory.py
delete mode 100644 tests/python/tir-transform/test_tir_transform_merge_static_shared_memory_allocations.py
delete mode 100644 web/tests/python/websock_rpc_test.py
(tvm) 01/01: [REFACTOR] Phase out te.schedule python components
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 6ce1fdb225837c60fc19e8f6113578983bd18f19
Author: tqchen <ti...@gmail.com>
AuthorDate: Sat Feb 15 11:21:50 2025 -0500
[REFACTOR] Phase out te.schedule python components
This PR phases out te.schedule python components.
te.compute is kept around for future usages.
tir.Schedule is a more modern version of the scheduling that we can use onwards.
Doing so also helps us to cleanup the testcases that relies on
explicit full build and execution. As we move future unit testcases
towards structural equality based unit tests.
---
apps/android_rpc/tests/android_rpc_test.py | 61 +-
apps/ios_rpc/tests/ios_rpc_test.py | 35 +-
docs/how_to/tutorials/cross_compilation_and_rpc.py | 18 +-
golang/sample/deploy.py | 56 -
jvm/README.md | 29 -
jvm/core/src/test/scripts/test_add_cpu.py | 43 -
jvm/core/src/test/scripts/test_add_gpu.py | 58 -
python/tvm/contrib/peak.py | 394 -------
python/tvm/contrib/sparse.py | 204 ----
python/tvm/contrib/tedd.py | 798 -------------
python/tvm/driver/build_module.py | 246 +---
python/tvm/exec/measure_peak.py | 52 -
python/tvm/relax/vm_build.py | 20 +-
python/tvm/te/__init__.py | 12 +-
python/tvm/te/autodiff.py | 67 --
python/tvm/te/hybrid/__init__.py | 101 --
python/tvm/te/hybrid/calls.py | 183 ---
python/tvm/te/hybrid/module.py | 113 --
python/tvm/te/hybrid/parser.py | 658 -----------
python/tvm/te/hybrid/preprocessor.py | 120 --
python/tvm/te/hybrid/runtime.py | 175 ---
python/tvm/te/hybrid/utils.py | 103 --
python/tvm/te/schedule.py | 665 -----------
python/tvm/te/tensor.py | 10 -
python/tvm/te/tensor_intrin.py | 146 ---
python/tvm/testing/utils.py | 3 +-
python/tvm/tir/buffer.py | 23 -
python/tvm/topi/__init__.py | 4 -
python/tvm/topi/argwhere.py | 197 ----
python/tvm/topi/nn/conv2d.py | 85 --
python/tvm/topi/random/__init__.py | 22 -
python/tvm/topi/random/kernel.py | 657 -----------
python/tvm/topi/sparse_fill_empty_rows.py | 109 --
python/tvm/topi/transform.py | 30 -
python/tvm/topi/unique.py | 234 ----
python/tvm/topi/vision/__init__.py | 25 -
python/tvm/topi/vision/nms.py | 1183 --------------------
python/tvm/topi/vision/nms_util.py | 338 ------
python/tvm/topi/vision/rcnn/__init__.py | 22 -
python/tvm/topi/vision/rcnn/proposal.py | 448 --------
python/tvm/topi/vision/rcnn/roi_align.py | 228 ----
python/tvm/topi/vision/rcnn/roi_pool.py | 95 --
python/tvm/topi/vision/reorg.py | 42 -
python/tvm/topi/vision/ssd/__init__.py | 22 -
python/tvm/topi/vision/ssd/multibox.py | 369 ------
python/tvm/utils/__init__.py | 19 -
python/tvm/utils/roofline/__init__.py | 279 -----
python/tvm/utils/roofline/cuda.py | 407 -------
python/tvm/utils/roofline/registry.py | 111 --
python/tvm/utils/roofline/x86.py | 331 ------
.../test_minimal_target_codegen_llvm.py | 22 +-
.../test_runtime_ndarray.py | 22 -
.../python/codegen/test_target_codegen_aarch64.py | 59 +-
tests/python/codegen/test_target_codegen_arm.py | 19 +-
tests/python/codegen/test_target_codegen_bool.py | 25 +-
tests/python/codegen/test_target_codegen_c_host.py | 118 +-
.../codegen/test_target_codegen_cross_llvm.py | 14 +-
tests/python/codegen/test_target_codegen_cuda.py | 539 +++------
tests/python/codegen/test_target_codegen_device.py | 49 +-
tests/python/codegen/test_target_codegen_extern.py | 25 +-
.../python/codegen/test_target_codegen_hexagon.py | 58 +-
tests/python/codegen/test_target_codegen_llvm.py | 433 +++----
tests/python/codegen/test_target_codegen_opencl.py | 9 +-
tests/python/codegen/test_target_codegen_rocm.py | 75 +-
tests/python/codegen/test_target_codegen_vulkan.py | 247 +---
tests/python/codegen/test_target_codegen_x86.py | 6 +-
tests/python/contrib/test_cblas.py | 16 +-
tests/python/contrib/test_dlpack.py | 3 +-
tests/python/contrib/test_gemm_acc16.py | 105 --
tests/python/contrib/test_gemm_acc32_vnni.py | 115 --
.../test_hexagon/conv2d/test_conv2d_blocked.py | 207 ----
.../test_hexagon/conv2d/test_conv2d_conv2d.py | 252 -----
.../test_hexagon/test_2d_physical_buffers.py | 41 +-
tests/python/contrib/test_hexagon/test_launcher.py | 12 +-
.../contrib/test_hexagon/test_maxpool2d_blocked.py | 158 ---
tests/python/contrib/test_hipblas.py | 6 +-
tests/python/contrib/test_miopen.py | 136 ---
tests/python/contrib/test_mps.py | 26 +-
tests/python/contrib/test_random.py | 9 +-
tests/python/contrib/test_rocblas.py | 6 +-
tests/python/contrib/test_sort.py | 41 +-
tests/python/contrib/test_sparse.py | 123 --
tests/python/runtime/test_runtime_dlpack.py | 4 +-
tests/python/runtime/test_runtime_measure.py | 3 +-
tests/python/runtime/test_runtime_module_export.py | 213 +---
tests/python/runtime/test_runtime_module_load.py | 16 +-
.../python/runtime/test_runtime_module_property.py | 12 +-
tests/python/runtime/test_runtime_rpc.py | 3 +-
tests/python/runtime/test_runtime_trace.py | 21 +-
tests/python/te/test_te_autodiff.py | 351 ------
tests/python/te/test_te_build_lower.py | 65 --
tests/python/te/test_te_group.py | 90 --
tests/python/te/test_te_hybrid_script.py | 872 ---------------
tests/python/te/test_te_schedule.py | 382 -------
.../python/te/test_te_schedule_bound_inference.py | 512 ---------
.../te/test_te_schedule_bound_inference_tiling.py | 62 -
tests/python/te/test_te_schedule_graph.py | 142 ---
tests/python/te/test_te_schedule_lstm.py | 91 --
tests/python/te/test_te_schedule_ops.py | 695 ------------
...te_schedule_postproc_rewrite_for_tensor_core.py | 231 ----
tests/python/te/test_te_schedule_tensor_core.py | 461 --------
tests/python/te/test_te_schedule_tensorize.py | 392 -------
tests/python/te/test_te_tensor.py | 202 +---
tests/python/te/test_te_transform_layout.py | 592 ----------
.../tir-analysis/test_tir_analysis_usedef.py | 36 -
.../test_tir_analysis_verify_gpu_code.py | 434 -------
.../test_tir_analysis_verify_memory.py | 121 --
tests/python/tir-base/test_lower_build.py | 17 -
tests/python/tir-base/test_tir_buffer.py | 79 --
tests/python/tir-base/test_tir_intrin.py | 68 +-
tests/python/tir-base/test_tir_ir_builder.py | 565 ----------
.../test_tir_transform_compact_buffer_region.py | 14 -
.../test_tir_transform_convert_blocks_to_opaque.py | 9 -
.../test_tir_transform_flatten_buffer.py | 13 -
.../tir-transform/test_tir_transform_hoist_if.py | 59 -
.../test_tir_transform_inject_copy_intrin.py | 124 --
.../test_tir_transform_inject_rolling_buffer.py | 177 ---
...test_tir_transform_instrument_bound_checkers.py | 608 ----------
.../test_tir_transform_loop_partition.py | 325 ------
...t_tir_transform_lower_cross_thread_reduction.py | 15 -
.../test_tir_transform_lower_init_block.py | 9 -
.../test_tir_transform_lower_intrin.py | 4 +-
.../test_tir_transform_lower_opaque_block.py | 9 -
.../test_tir_transform_lower_warp_memory.py | 356 ------
.../test_tir_transform_make_packed_api.py | 26 -
...form_merge_dynamic_shared_memory_allocations.py | 304 -----
...sform_merge_static_shared_memory_allocations.py | 203 ----
.../test_tir_transform_narrow_datatype.py | 41 -
...sform_plan_update_buffer_allocation_location.py | 11 -
.../tir-transform/test_tir_transform_simplify.py | 64 --
.../test_tir_transform_split_host_device.py | 39 -
.../test_tir_transform_storage_flatten.py | 63 --
.../test_tir_transform_storage_rewrite.py | 320 ------
.../test_tir_transform_thread_sync.py | 61 -
.../test_tir_transform_unify_thread_binding.py | 11 -
.../test_tir_transform_unroll_loop.py | 18 -
.../tir-transform/test_tir_transform_vectorize.py | 25 +-
web/tests/python/webgpu_rpc_test.py | 3 +-
web/tests/python/websock_rpc_test.py | 91 --
139 files changed, 773 insertions(+), 21587 deletions(-)
diff --git a/apps/android_rpc/tests/android_rpc_test.py b/apps/android_rpc/tests/android_rpc_test.py
index 0027cc4ba3..6cdef0170e 100644
--- a/apps/android_rpc/tests/android_rpc_test.py
+++ b/apps/android_rpc/tests/android_rpc_test.py
@@ -56,37 +56,17 @@ def test_rpc_module():
tracker = rpc.connect_tracker(tracker_host, tracker_port)
remote = tracker.request(key, priority=0, session_timeout=60)
- # Compile the Graph for CPU target
- s = te.create_schedule(B.op)
- xo, xi = s[B].split(B.op.axis[0], factor=64)
- s[B].parallel(xi)
- s[B].pragma(xo, "parallel_launch_point")
- s[B].pragma(xi, "parallel_barrier_when_finish")
- f = tvm.build(s, [A, B], target, name="myadd_cpu")
- path_dso_cpu = temp.relpath("cpu_lib.so")
- f.export_library(path_dso_cpu, fcompile=ndk.create_shared)
+ mod = tvm.IRModule.from_expr(
+ te.create_prim_func([A, B]).with_attr("global_symbol", "myadd")
+ )
+ sch = tvm.tir.Schedule(mod)
+ (x,) = sch.get_loops(block=sch.get_block("B"))
+ xo, xi = sch.split(i, [None, 32])
+ sch.bind(xo, "blockIdx.x")
+ sch.bind(xi, "threadIdx.x")
- # Execute the portable graph on cpu target
- print("Run CPU test ...")
- dev = remote.cpu(0)
- remote.upload(path_dso_cpu)
- f2 = remote.load_module("cpu_lib.so")
- a = tvm.nd.array(a_np, dev)
- b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
- time_f = f2.time_evaluator(f2.entry_name, dev, number=10)
- cost = time_f(a, b).mean
- print("%g secs/op\n" % cost)
- np.testing.assert_equal(b.numpy(), a.numpy() + 1)
-
- # Compile the Graph for OpenCL target
if test_opencl:
- s = te.create_schedule(B.op)
- xo, xi = s[B].split(B.op.axis[0], factor=64)
- s[B].bind(xi, te.thread_axis("threadIdx.x"))
- s[B].bind(xo, te.thread_axis("blockIdx.x"))
- # Build the dynamic lib.
- # If we don't want to do metal and only use cpu, just set target to be target
- f = tvm.build(s, [A, B], tvm.target.Target("opencl", host=target), name="myadd")
+ f = tvm.build(sch.mod, target=tvm.target.Target("opencl", host=target))
path_dso_cl = temp.relpath("dev_lib_cl.so")
f.export_library(path_dso_cl, fcompile=ndk.create_shared)
@@ -101,29 +81,6 @@ def test_rpc_module():
print("%g secs/op\n" % cost)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
- # Compile the Graph for Vulkan target
- if test_vulkan:
- s = te.create_schedule(B.op)
- xo, xi = s[B].split(B.op.axis[0], factor=64)
- s[B].bind(xi, te.thread_axis("threadIdx.x"))
- s[B].bind(xo, te.thread_axis("blockIdx.x"))
- # Build the dynamic lib.
- # If we don't want to do metal and only use cpu, just set target to be target
- f = tvm.build(s, [A, B], tvm.target.Target("vulkan", host=target), name="myadd")
- path_dso_vulkan = temp.relpath("dev_lib_vulkan.so")
- f.export_library(path_dso_vulkan, fcompile=ndk.create_shared)
-
- print("Run GPU(Vulkan Flavor) test ...")
- dev = remote.vulkan(0)
- remote.upload(path_dso_vulkan)
- f1 = remote.load_module("dev_lib_vulkan.so")
- a = tvm.nd.array(a_np, dev)
- b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
- time_f = f1.time_evaluator(f1.entry_name, dev, number=10)
- cost = time_f(a, b).mean
- print("%g secs/op\n" % cost)
- np.testing.assert_equal(b.numpy(), a.numpy() + 1)
-
if __name__ == "__main__":
test_rpc_module()
diff --git a/apps/ios_rpc/tests/ios_rpc_test.py b/apps/ios_rpc/tests/ios_rpc_test.py
index f0c31cd7d2..00e70bd010 100644
--- a/apps/ios_rpc/tests/ios_rpc_test.py
+++ b/apps/ios_rpc/tests/ios_rpc_test.py
@@ -50,25 +50,21 @@ def test_rpc_module(host, port, key, mode):
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
temp = utils.tempdir()
- s = te.create_schedule(B.op)
- xo, xi = s[B].split(B.op.axis[0], factor=64)
- s[B].bind(xi, te.thread_axis("threadIdx.x"))
- s[B].bind(xo, te.thread_axis("blockIdx.x"))
+ mod = tvm.IRModule.from_expr(
+ te.create_prim_func([A, B]).with_attr("global_symbol", "myadd")
+ )
+ sch = tvm.tir.Schedule(mod)
+ (i,) = sch.get_loops(block=sch.get_block("B"))
+ i0, i1 = sch.split(i, [None, 32])
+ sch.bind(i0, "blockIdx.x")
+ sch.bind(i1, "threadIdx.x")
+
# Build the dynamic lib.
# If we don't want to do metal and only use cpu, just set target to be target
- f = tvm.build(s, [A, B], tvm.target.Target("metal", host=target), name="myadd")
+ f = tvm.build(sch.mod, target=tvm.target.Target("metal", host=target))
path_dso1 = temp.relpath("dev_lib.dylib")
f.export_library(path_dso1, fcompile=xcode.create_dylib, arch=arch, sdk=sdk)
- s = te.create_schedule(B.op)
- xo, xi = s[B].split(B.op.axis[0], factor=64)
- s[B].parallel(xi)
- s[B].pragma(xo, "parallel_launch_point")
- s[B].pragma(xi, "parallel_barrier_when_finish")
- f = tvm.build(s, [A, B], target, name="myadd_cpu")
- path_dso2 = temp.relpath("cpu_lib.dylib")
- f.export_library(path_dso2, fcompile=xcode.create_dylib, arch=arch, sdk=sdk)
-
# connect to the proxy
if mode == "tracker":
remote = MODES[mode](host, port).request(key)
@@ -84,17 +80,6 @@ def test_rpc_module(host, port, key, mode):
cost = time_f(a, b).mean
print("Metal: %g secs/op" % cost)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
- # CPU
- dev = remote.cpu(0)
- remote.upload(path_dso2)
- f2 = remote.load_module("cpu_lib.dylib")
- a_np = np.random.uniform(size=1024).astype(A.dtype)
- a = tvm.nd.array(a_np, dev)
- b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), dev)
- time_f = f2.time_evaluator(f2.entry_name, dev, number=10)
- cost = time_f(a, b).mean
- print("CPU: %g secs/op" % cost)
- np.testing.assert_equal(b.numpy(), a.numpy() + 1)
if __name__ == "__main__":
diff --git a/docs/how_to/tutorials/cross_compilation_and_rpc.py b/docs/how_to/tutorials/cross_compilation_and_rpc.py
index c7e302693d..674a4588ed 100644
--- a/docs/how_to/tutorials/cross_compilation_and_rpc.py
+++ b/docs/how_to/tutorials/cross_compilation_and_rpc.py
@@ -104,7 +104,9 @@ from tvm.contrib import utils
n = tvm.runtime.convert(1024)
A = te.placeholder((n,), name="A")
B = te.compute((n,), lambda i: A[i] + 1.0, name="B")
-s = te.create_schedule(B.op)
+mod = tvm.IRModule.from_expr(
+ te.create_prim_func([A, B]).with_attr("global_symbol", "add_one")
+)
######################################################################
# Then we cross compile the kernel.
@@ -119,7 +121,7 @@ if local_demo:
else:
target = "llvm -mtriple=armv7l-linux-gnueabihf"
-func = tvm.build(s, [A, B], target=target, name="add_one")
+func = tvm.build(mod, target=target, name="add_one")
# save the lib at a local temp folder
temp = utils.tempdir()
path = temp.relpath("lib.tar")
@@ -231,11 +233,13 @@ def run_opencl():
target = tvm.target.Target("opencl", host="llvm -mtriple=aarch64-linux-gnu")
# create schedule for the above "add one" compute declaration
- s = te.create_schedule(B.op)
- xo, xi = s[B].split(B.op.axis[0], factor=32)
- s[B].bind(xo, te.thread_axis("blockIdx.x"))
- s[B].bind(xi, te.thread_axis("threadIdx.x"))
- func = tvm.build(s, [A, B], target=target)
+ mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]))
+ sch = tvm.tir.Schedule(mod)
+ (x,) = sch.get_loops(block=sch.get_block("B"))
+ xo, xi = sch.split(i, [None, 32])
+ sch.bind(x, "blockIdx.x")
+ sch.bind(x, "threadIdx.x")
+ func = tvm.build(sch.mod, target=target)
remote = rpc.connect(opencl_device_host, opencl_device_port)
diff --git a/golang/sample/deploy.py b/golang/sample/deploy.py
deleted file mode 100644
index fa4cbd4335..0000000000
--- a/golang/sample/deploy.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# 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.
-"""
-Get Started with TVM Go
-=======================
-"""
-from __future__ import absolute_import, print_function
-
-import tvm
-from tvm import te
-import numpy as np
-
-# Global declarations of environment.
-
-tgt = "llvm"
-
-######################################################################
-# Describe the Computation
-# ------------------------
-n = te.var("n")
-A = te.placeholder((n,), name="A")
-B = te.placeholder((n,), name="B")
-C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
-
-######################################################################
-# Schedule the Computation
-# ------------------------
-s = te.create_schedule(C.op)
-
-######################################################################
-# Compilation
-# -----------
-fadd = tvm.build(s, [A, B, C], tgt, name="myadd")
-
-######################################################################
-# Save Compiled Module
-# --------------------
-from tvm.contrib import cc
-from tvm.contrib import utils
-
-fadd.save("deploy.o")
-cc.create_shared("deploy.so", ["deploy.o"])
diff --git a/jvm/README.md b/jvm/README.md
index c7535f0311..62b685010c 100644
--- a/jvm/README.md
+++ b/jvm/README.md
@@ -89,35 +89,6 @@ It is your job to verify the types of callback arguments, as well as the type of
You can register the Java function by `Function.register` and use `Function.getFunction` to get the registered function later.
-## Use TVM to Generate Shared Library
-
-There's nothing special for this part. The following Python snippet generate add_cpu.so which add two vectors on CPU.
-
-```python
-import os
-import tvm
-from tvm import te
-from tvm.contrib import cc, utils
-
-def test_add(target_dir):
- n = te.var("n")
- A = te.placeholder((n,), name='A')
- B = te.placeholder((n,), name='B')
- C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
- s = te.create_schedule(C.op)
- fadd = tvm.build(s, [A, B, C], "llvm", name="myadd")
-
- fadd.save(os.path.join(target_dir, "add_cpu.o"))
- cc.create_shared(os.path.join(target_dir, "add_cpu.so"),
- [os.path.join(target_dir, "add_cpu.o")])
-
-if __name__ == "__main__":
- import sys
- if len(sys.argv) != 2:
- sys.exit(-1)
- test_add(sys.argv[1])
-```
-
## Run the Generated Shared Library
The following code snippet demonstrate how to load generated shared library (add_cpu.so).
diff --git a/jvm/core/src/test/scripts/test_add_cpu.py b/jvm/core/src/test/scripts/test_add_cpu.py
deleted file mode 100644
index 9a93d4e746..0000000000
--- a/jvm/core/src/test/scripts/test_add_cpu.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# 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.
-import os
-
-import tvm
-from tvm import te
-from tvm.contrib import cc, utils
-
-
-def test_add(target_dir):
- n = te.var("n")
- A = te.placeholder((n,), name="A")
- B = te.placeholder((n,), name="B")
- C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
- s = te.create_schedule(C.op)
- fadd = tvm.build(s, [A, B, C], "llvm", name="myadd")
-
- fadd.save(os.path.join(target_dir, "add_cpu.o"))
- cc.create_shared(
- os.path.join(target_dir, "add_cpu.so"), [os.path.join(target_dir, "add_cpu.o")]
- )
-
-
-if __name__ == "__main__":
- import sys
-
- if len(sys.argv) != 2:
- sys.exit(-1)
- test_add(sys.argv[1])
diff --git a/jvm/core/src/test/scripts/test_add_gpu.py b/jvm/core/src/test/scripts/test_add_gpu.py
deleted file mode 100644
index 0eea5671ba..0000000000
--- a/jvm/core/src/test/scripts/test_add_gpu.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# 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.
-import os
-
-import tvm
-from tvm import te
-from tvm.contrib import cc, nvcc, utils
-
-
-@tvm.register_func("tvm_callback_cuda_compile", override=True)
-def tvm_callback_cuda_compile(code, target):
- ptx = nvcc.compile_cuda(code, target_format="ptx")
- return ptx
-
-
-def test_add(target_dir):
- if not tvm.runtime.enabled("cuda"):
- print("skip %s because cuda is not enabled..." % __file__)
- return
- n = te.var("n")
- A = te.placeholder((n,), name="A")
- B = te.placeholder((n,), name="B")
- C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
-
- s = te.create_schedule(C.op)
-
- bx, tx = s[C].split(C.op.axis[0], factor=64)
- s[C].bind(bx, te.thread_axis("blockIdx.x"))
- s[C].bind(tx, te.thread_axis("threadIdx.x"))
- fadd_cuda = tvm.build(s, [A, B, C], tvm.target.Target("cuda", host="llvm"), name="myadd")
-
- fadd_cuda.save(os.path.join(target_dir, "add_cuda.o"))
- fadd_cuda.imported_modules[0].save(os.path.join(target_dir, "add_cuda.ptx"))
- cc.create_shared(
- os.path.join(target_dir, "add_cuda.so"), [os.path.join(target_dir, "add_cuda.o")]
- )
-
-
-if __name__ == "__main__":
- import sys
-
- if len(sys.argv) != 2:
- sys.exit(-1)
- test_add(sys.argv[1])
diff --git a/python/tvm/contrib/peak.py b/python/tvm/contrib/peak.py
deleted file mode 100644
index 78dae846d6..0000000000
--- a/python/tvm/contrib/peak.py
+++ /dev/null
@@ -1,394 +0,0 @@
-# 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.
-# pylint: disable=invalid-name
-"""measure bandwidth and compute peak"""
-
-import logging
-import tvm
-from tvm import te
-from tvm.target import Target
-from . import utils
-from .. import rpc
-
-
-def _convert_to_remote(func, remote):
- """convert module function to remote rpc function"""
- temp = utils.tempdir()
- path_dso = temp.relpath("tmp_func.tar")
- func.export_library(path_dso)
-
- remote.upload(path_dso)
- func = remote.load_module("tmp_func.tar")
- return func
-
-
-def measure_bandwidth_sum(
- total_item,
- item_per_thread,
- stride,
- base_type,
- bits,
- lanes,
- target,
- target_host,
- remote,
- dev,
- n_times,
-):
- """measure memory bandwidth of gpu by product reduction for a given type
-
- The IR for measurement is
-
- for each thread
- for i in 1..num_per_thread:
- y[global_id] = y[global_id] * x[base + i * stride]
-
- Parameters
- ----------
- total_item: int
- number of elements in input array
- item_per_thread: int
- number of elements each thread accumulates
- stride: int
- stride in memory access
- base_type: str
- can be "int", "float"
- bits: int
- can be 16, 32
- lanes: int
- lane of the vector type, can be 1, 2, 4, 8, 16
- target: :any:`tvm.target.Target`
- the target and option of the compilation.
- target_host : str or :any:`tvm.target.Target`
- host compilation target
- dev: Device
- the device of array
- remote: tvm.rpc.RPCSession
- remote rpc session
- n_times: int
- number of runs for taking mean
-
- Returns
- -------
- GBPS: float
- gigabyte per second
- """
- target, target_host = Target.canon_target_and_host(target, target_host)
-
- n, m = total_item, item_per_thread
- n //= lanes
-
- base_type = str(base_type) + str(bits)
- dtype = base_type if lanes == 1 else base_type + "x" + str(lanes)
-
- k = te.reduce_axis((0, m), name="k")
-
- x = te.placeholder((n,), dtype=dtype, name="x")
- op = te.comm_reducer(lambda x, y: x * y, lambda t: tvm.tir.const(1, dtype=t), name="sum")
- y = te.compute(
- (n // m,), lambda i: op(x[i // stride * stride * m + i % stride + k * stride], axis=k)
- )
- s = te.create_schedule(y.op)
-
- yo, yi = s[y].split(y.op.axis[0], target.max_num_threads)
- s[y].bind(yo, te.thread_axis("blockIdx.x"))
- s[y].bind(yi, te.thread_axis("threadIdx.x"))
- s[y].unroll(k)
-
- try:
- func = tvm.build(s, [x, y], target)
-
- x = tvm.nd.empty((n,), dtype=dtype, device=dev)
- y = tvm.nd.empty((n // m,), dtype=dtype, device=dev)
-
- func = _convert_to_remote(func, remote)
- time_f = func.time_evaluator(func.entry_name, dev, number=n_times)
- time = time_f(x, y).mean
- except tvm._ffi.base.TVMError:
- # build error (occur when device does not support half)
- return -1
-
- return 1.0 * (total_item * bits / 8) / 1e9 / time
-
-
-def measure_bandwidth_all_types(
- total_item, item_per_thread, n_times, target, target_host, remote, dev, verbose=True
-):
- """measure memory bandwidth for all types
-
- Parameters
- ----------
- total_item: int
- number of elements in input array
- item_per_thread: int
- number of elements each thread accmulates
- n_times: int
- number of runs for averaging
- target: :any:`tvm.target.Target`
- the target and option of the compilation.
- target_host : str or :any:`tvm.target.Target`
- host compilation target
- remote: tvm.rpc.RPCSession
- remote rpc session
- dev: Device
- the device of array
- verbose: bool
- whether outputs immediate result
-
- Returns
- -------
- result: list
- a list of (type_name, GBPS) pairs
- """
- target, target_host = Target.canon_target_and_host(target, target_host)
- max_threads = target.max_num_threads
-
- result = []
- for base_type in ["float"]:
- for bits in [32]:
- for lanes in [1, 2, 4, 8, 16]:
- max_speed = -1e9
- # try different strides
- for stride in [max_threads, total_item // (lanes * item_per_thread)]:
- speed = measure_bandwidth_sum(
- total_item,
- item_per_thread,
- stride,
- base_type,
- bits,
- lanes,
- target,
- target_host,
- remote,
- dev,
- n_times,
- )
- max_speed = max(max_speed, speed)
- type_name = base_type + str(bits)
- result.append([f"{type_name}x{lanes}", max_speed])
- if verbose:
- logging.info("\t%-10s %.2f GBPS", result[-1][0], result[-1][1])
- return result
-
-
-def measure_compute_mad(
- total_item, item_per_thread, base_type, bits, lanes, target, target_host, remote, dev, n_times
-):
- """measure peak compute speed by computing mad for a type
-
- The IR for measurement is
-
- for each thread
- for i in 1..item_per_thread
- x = mad(x, x, y)
- y = mad(y, y, x)
-
- Parameters
- ----------
- total_item: int
- number of elements in input array
- item_per_thread: int
- number of operations each thread does
- base_type: str
- can be "int", "float"
- bits: int
- can be 16, 32
- lanes: int
- lane of the vector type, can be 1, 2, 4, 8, 16
- target: :any:`tvm.target.Target`
- the target and option of the compilation.
- target_host : str or :any:`tvm.target.Target`
- host compilation target
- remote: tvm.rpc.RPCSession
- if it is not None, use remote rpc session
- dev: Device
- the device of array
- n_times: int
- number of runs for taking mean
-
- Returns
- -------
- GOPS: float
- giga operation per second
- """
- target, target_host = Target.canon_target_and_host(target, target_host)
-
- n = total_item
-
- if bits >= 64 or lanes >= 16:
- n //= 2
-
- max_threads = target.max_num_threads
-
- base_type = str(base_type) + str(bits)
- dtype = base_type if lanes == 1 else base_type + "x" + str(lanes)
-
- def extern(ins, outs):
- # pylint: disable=unused-argument
- """construct measurement function by building IR directly"""
- ib = tvm.tir.ir_builder.create()
-
- bx = te.thread_axis("blockIdx.x")
- tx = te.thread_axis("threadIdx.x")
-
- ib.scope_attr(bx, "thread_extent", n // max_threads)
- ib.scope_attr(tx, "thread_extent", max_threads)
-
- idx = bx.var * max_threads + tx.var
-
- a = ib.allocate(dtype, (1), name="a", scope="local")
- b = ib.allocate(dtype, (1), name="b", scope="local")
-
- a[0] = outs[0].vload(idx, dtype)
- b[0] = outs[0].vload(idx, dtype)
-
- if base_type.find("float") != -1:
-
- def mad_func(x, y):
- return x * x + y
-
- else:
-
- def mad_func(x, y):
- return y * y + x
-
- for _ in range(item_per_thread // 4 // lanes):
- a[0] = mad_func(a[0], b[0])
- b[0] = mad_func(b[0], a[0])
-
- ib.emit(outs[0].vstore(idx, b[0]))
- return ib.get()
-
- y = te.extern((n,), [], extern, name="y", dtype=dtype)
- s = te.create_schedule(y.op)
-
- try:
- func = tvm.build(s, [y], target)
- func = _convert_to_remote(func, remote)
- time_f = func.time_evaluator(func.entry_name, dev, number=n_times)
- y = tvm.nd.empty((n,), dtype=dtype, device=dev)
- time = time_f(y).mean
- except tvm._ffi.base.TVMError:
- # build error (occur when device does not support half)
- return -1
-
- return 1.0 * (n * item_per_thread) / 1e9 / time
-
-
-def measure_compute_all_types(
- total_item, item_per_thread, n_times, target, target_host, remote, dev, verbose=True
-):
- """measure peak flops for all types
-
- Parameters
- ----------
- total_item: int
- number of elements in input array
- item_per_thread: int
- number of elements each thread accmulates
- n_times: int
- number of runs for averaging
- target: :any:`tvm.target.Target`
- the target and option of the compilation.
- target_host : str or :any:`tvm.target.Target`
- host compilation target
- remote: tvm.rpc.RPCSession
- remote rpc session
- dev: Device
- the device of array
- verbose: bool
- whether outputs immediate result
-
- Returns
- -------
- result: list
- a list of (type_name, GFLOPS/GIOPS) pairs
- """
- target, target_host = Target.canon_target_and_host(target, target_host)
-
- result = []
- for base_type in ["float", "int"]:
- for bits in [16, 32, 64]:
- for lanes in [1, 2, 4, 8, 16]:
- if base_type == "int" and bits != 32: # only measure int32
- continue
-
- max_speed = -1e9
- for per_thread in [item_per_thread // 2, item_per_thread, item_per_thread * 2]:
- speed = measure_compute_mad(
- total_item,
- per_thread,
- base_type,
- bits,
- lanes,
- target,
- target_host,
- remote,
- dev,
- n_times,
- )
- max_speed = max(max_speed, speed)
- type_name = base_type + str(bits)
- result.append([f"{type_name}x{lanes}", max_speed])
-
- unit = "GFLOPS" if base_type == "float" else "GIOPS"
-
- if verbose:
- logging.info("\t%-10s %.2f %s", result[-1][0], result[-1][1], unit)
-
- return result
-
-
-def measure_peak_all(target, target_host, host, port):
- """measure memory bandwidth and peak compute for gpu devices
-
- Parameters
- ----------
- target: str or :any:`tvm.target.Target`
- target_host: str
- host: str
- port: int
- """
-
- target, target_host = Target.canon_target_and_host(target, target_host)
- remote = rpc.connect(host, port)
- n_times = 20
-
- bandwidth_total_item = 1 << 25
- bandwidth_item_per_thread = 32
-
- compute_total_item = 1 << 21
- compute_item_per_thread = 4096
-
- if str(target).startswith("opencl"):
- dev = remote.cl()
- elif str(target).startswith("cuda"):
- dev = remote.cuda()
- elif str(target).startswith("metal"):
- dev = remote.metal()
- else:
- raise RuntimeError("Unsupported target")
-
- logging.info("========== measure memory bandwidth ==========")
- measure_bandwidth_all_types(
- bandwidth_total_item, bandwidth_item_per_thread, n_times, target, target_host, remote, dev
- )
-
- logging.info("========== measure peak compute ==========")
- measure_compute_all_types(
- compute_total_item, compute_item_per_thread, n_times, target, target_host, remote, dev
- )
diff --git a/python/tvm/contrib/sparse.py b/python/tvm/contrib/sparse.py
deleted file mode 100644
index 9f94ff24f9..0000000000
--- a/python/tvm/contrib/sparse.py
+++ /dev/null
@@ -1,204 +0,0 @@
-# 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.
-"""Tensor and Operation class for computation declaration."""
-# pylint: disable=invalid-name
-import warnings
-import numpy as _np
-from tvm.runtime import ndarray as _nd
-from tvm import te
-from tvm.tir import expr as _expr
-from tvm.te import tensor as _tensor
-
-
-float32 = "float32"
-itype = "int32"
-
-
-class CSRNDArray(object):
- """Sparse tensor object in CSR format."""
-
- def __init__(self, arg1, device=None, shape=None):
- """Construct a sparse matrix in CSR format.
-
- Parameters
- ----------
- arg1 : numpy.ndarray or a tuple with (data, indices, indptr)
- The corresponding a dense numpy array,
- or a tuple for constructing a sparse matrix directly.
-
- device: Device
- The corresponding device.
-
- shape : tuple of int
- The shape of the array
- """
- if isinstance(arg1, tuple):
- assert len(arg1) == 3
- self.data, self.indices, self.indptr = arg1
- self.shape = shape
- elif isinstance(arg1, _np.ndarray):
- source_array = arg1
- ridx, cidx = _np.nonzero(source_array)
- data = source_array[ridx, cidx]
- self.data = _nd.array(data, device)
- indices = _np.nonzero(source_array)[1].astype(itype)
- self.indices = _nd.array(indices, device)
- indptr = [0] + _np.apply_along_axis(
- _np.count_nonzero, axis=1, arr=source_array
- ).tolist()
- indptr = _np.cumsum(_np.array(indptr, itype)).astype(itype)
- self.indptr = _nd.array(indptr, device)
- self.shape = source_array.shape
- else:
- raise RuntimeError(
- f"Construct CSRNDArray with either a tuple (data, indices, indptr) "
- f"or a numpy.array, can't handle type {type(arg1)}."
- )
- self.stype = "csr"
- self.dtype = self.data.dtype
- assert self.shape is not None
- assert isinstance(self.data, _nd.NDArray)
- assert isinstance(self.indices, _nd.NDArray)
- assert str(self.indices.dtype) == "int32" or str(self.indices.dtype) == "int64", str(
- self.indices.dtype
- )
- assert isinstance(self.indptr, _nd.NDArray)
- assert str(self.indptr.dtype) == "int32" or str(self.indptr.dtype) == "int64", str(
- self.indptr.dtype
- )
-
- def asnumpy(self):
- """Construct a full matrix and convert it to numpy array. This API will be deprecated
- in TVM v0.8 release. Please use `numpy` instead."""
- warnings.warn(
- "CSRNDArray.asnumpy() will be deprecated in TVM v0.8 release. "
- "Please use CSRNDArray.numpy() instead.",
- DeprecationWarning,
- )
- return self.numpy()
-
- def numpy(self):
- """Construct a full matrix and convert it to numpy array."""
- full = _np.zeros(self.shape, self.dtype)
- ridx = _np.diff(self.indptr.numpy())
- ridx = _np.hstack([_np.ones((v,), itype) * i for i, v in enumerate(ridx)])
- full[ridx, self.indices.numpy().astype(itype)] = self.data.numpy()
- return full
-
-
-def array(source_array, device=None, shape=None, stype="csr"):
- """Construct a sparse NDArray from numpy.ndarray"""
- ret = None
- if stype == "csr":
- ret = CSRNDArray(source_array, shape=shape, device=device)
- else:
- raise NotImplementedError(f"stype={stype} is not supported yet.")
- return ret
-
-
-class SparsePlaceholderOp(object):
- """Placeholder class for sparse tensor representations."""
-
- def __init__(self, shape, nonzeros, dtype, name):
- # pylint: disable=unused-argument
- """Contructing a bare bone structure for a sparse matrix
-
- Parameters
- ----------
- shape: Tuple of Expr
- The shape of the tensor
-
- nonzeros: int
- The number of non-zero values
-
- dtype: str, optional
- The data type of the tensor
-
- name: str, optional
- The name hint of the tensor
- """
- self.shape = shape
- self.dtype = dtype
- self.name = name
- self.stype = "unknown"
-
-
-class CSRPlaceholderOp(SparsePlaceholderOp):
- """Placeholder class for CSR based sparse tensor representation."""
-
- def __init__(self, shape, nonzeros, dtype, name):
- """Contructing a bare bone structure for a csr_matrix
-
- Parameters
- ----------
- shape: Tuple of Expr
- The shape of the tensor
-
- nonzeros: int
- The number of non-zero values
-
- dtype: str, optional
- The data type of the tensor
-
- name: str, optional
- The name hint of the tensor
- """
- SparsePlaceholderOp.__init__(self, shape, nonzeros, dtype, name)
- self.stype = "csr"
- self.data = te.placeholder((nonzeros,), dtype=dtype, name=self.name + "_data")
- self.indices = te.placeholder((nonzeros,), dtype=itype, name=self.name + "_indices")
- self.indptr = te.placeholder((self.shape[0] + 1,), dtype=itype, name=self.name + "_indptr")
- assert isinstance(self.data, _tensor.Tensor)
- assert isinstance(self.indices, _tensor.Tensor)
- assert isinstance(self.indptr, _tensor.Tensor)
-
-
-def placeholder(shape, nonzeros=None, dtype=None, name="placeholder", stype=None):
- """Construct an empty sparse tensor object.
-
- Parameters
- ----------
- shape: Tuple of Expr
- The shape of the tensor
-
- nonzeros: int
- The number of non-zero values
-
- dtype: str, optional
- The data type of the tensor
-
- name: str, optional
- The name hint of the tensor
-
- stype: str, optional
- The name storage type of the sparse tensor (e.g. csr, coo, ell)
-
- Returns
- -------
- tensor: SparsePlaceholderOp
- The created sparse tensor placeholder
- """
- shape = (shape,) if isinstance(shape, _expr.PrimExpr) else shape
- nonzeros = 0 if nonzeros is None else nonzeros
- dtype = float32 if dtype is None else dtype
- stype = "csr" if stype is None else stype
- ret = None
- if stype == "csr":
- ret = CSRPlaceholderOp(shape=shape, nonzeros=nonzeros, dtype=dtype, name=name)
- else:
- raise NotImplementedError(f"stype={stype} is not supported yet.")
- return ret
diff --git a/python/tvm/contrib/tedd.py b/python/tvm/contrib/tedd.py
deleted file mode 100644
index 6802977297..0000000000
--- a/python/tvm/contrib/tedd.py
+++ /dev/null
@@ -1,798 +0,0 @@
-# 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.
-# pylint: disable=import-outside-toplevel, nested-min-max
-"""Tensor Expression Debug Display (TEDD), visualizing Tensor Expression"""
-import html
-import json
-import warnings
-from graphviz import Digraph
-from graphviz import Source
-import tvm
-
-TVMDD_TABLE_BODY_WIDTH = 30
-# Must match enum IterVarType defined in include/tvm/expr.h
-ITERVAR_TYPE_STRING_MAP = {
- 0: ("kDataPar", "#FFFFFF"),
- 1: ("kThreadIndex", "#2980B9"),
- 2: ("kCommReduce", "#FAD7A0"),
- 3: ("kOrdered", "#D35400"),
- 4: ("kOpaque", "#ABB2B9"),
- 5: ("kUnrolled", "#D2B4DE"),
- 6: ("kVectorized", "#AED6F1"),
- 7: ("kParallelized", "#F5B7B1"),
- 8: ("kTensorized", "#A9DFBF"),
-}
-
-PALETTE = {
- 0: "#000000",
- 1: "#922B21",
- 2: "#76448A",
- 3: "#1F618D",
- 4: "#148F77",
- 5: "#B7950B",
- 6: "#AF601A",
- 7: "#F5B7B1",
- 8: "#A9DFBF",
-}
-
-PALETTE_SIZE = 9
-
-
-def dom_path_to_string(dom_path, prefix=""):
- path_string = prefix
- for index in dom_path:
- path_string = path_string + "_" + str(index)
- return path_string
-
-
-def insert_dot_id(sch):
- """Insert unique ID for each node in the DOM tree.
- They are used as Dot node ID.
- """
- for stage_idx, stage in enumerate(sch["stages"]):
- dom_path = [stage_idx]
- stage["id"] = dom_path_to_string(dom_path, stage["type"])
- for itervar_idx, itervar in enumerate(stage["all_itervars"]):
- dom_path = [stage_idx, itervar_idx]
- itervar["id"] = dom_path_to_string(dom_path, itervar["type"])
- for rel_idx, rel in enumerate(stage["relations"]):
- dom_path = [stage_idx, rel_idx]
- rel["id"] = dom_path_to_string(dom_path, rel["type"])
- for tensor_idx, tensor in enumerate(stage["output_tensors"]):
- dom_path = [stage_idx, tensor_idx]
- tensor["id"] = dom_path_to_string(dom_path, tensor["type"])
- return sch
-
-
-def itervar_equal(iv_a, iv_b):
- """A helper method that compares the equality of two iterative variables"""
- # Adopt the following method to assure the equality between two itervars.
- # The plain comparison might fail (i.e. iv_a == iv_b) after the change of
- # domain bounds from InferBound.
- def _var_equal(v_a, v_b):
- condtions = [
- v_a.name == v_b.name,
- v_a.dtype == v_b.dtype,
- v_a.type_annotation == v_b.type_annotation,
- ]
- return all(c for c in condtions)
-
- condtions = [
- _var_equal(iv_a.var, iv_b.var),
- iv_a.iter_type == iv_b.iter_type,
- iv_a.thread_tag == iv_b.thread_tag,
- ]
- return all(c for c in condtions)
-
-
-class ObjectManager:
- """A helper class tracking schedule objects, e.g. stage, IterVar,
- relationship, and tensor, to their DOM path."""
-
- def __init__(self, sch):
- self.dict = {}
- for stage_idx, stage in enumerate(sch.stages):
- self.dict[stage] = [stage_idx]
- for itervar_idx, itervar in enumerate(stage.all_iter_vars):
- self.dict[itervar] = [stage_idx, itervar_idx]
- # the itervars of leaf should also be mapped to the original one
- for leaf_iv in stage.leaf_iter_vars:
- if itervar_equal(leaf_iv, itervar):
- self.dict[leaf_iv] = [stage_idx, itervar_idx]
- for rel_idx, rel in enumerate(stage.relations):
- self.dict[rel] = [stage_idx, rel_idx]
- for tensor_idx in range(stage.op.num_outputs):
- self.dict[frozenset({stage.op.name, tensor_idx})] = [stage_idx, tensor_idx]
-
- def get_dom_path(self, obj):
- if obj is None:
- return None
- assert obj in self.dict, "Node is no found."
- return self.dict[obj]
-
-
-def get_or_create_dot_id(obj, prefix="", assert_on_missing=False):
- """If obj's ID has been registered, return it.
- If not, either assert or create a unique and legal ID, register and
- return it, according to assert_on_missing.
- ID must be a unique and legal Dotty ID.
-
- Parameters
- ----------
- obj : objet
- Serve as the key to the ID.
-
- prefix : string
- Prefix to attach to the ID. Usually use obj's non-unique
- name as prefix.
-
- assert_on_missing : bool
- Assert or not if object doesn't have a registered ID.
- """
- prefix = prefix.replace(".", "_")
- if not hasattr(get_or_create_dot_id, "obj_id_dict"):
- get_or_create_dot_id.obj_id_dict = {}
- if obj not in get_or_create_dot_id.obj_id_dict:
- if assert_on_missing:
- assert False, "dot_id " + str(obj) + " has not been registered."
- else:
- get_or_create_dot_id.obj_id_dict[obj] = prefix + hex(id(obj))
- return get_or_create_dot_id.obj_id_dict[obj]
-
-
-def get_port_id(is_input, index):
- return "I_" + str(index) if is_input else "O_" + str(index)
-
-
-def get_itervar_type_info(iter_type):
- assert iter_type < len(ITERVAR_TYPE_STRING_MAP), "Unknown IterVar type: " + str(iter_type)
- return ITERVAR_TYPE_STRING_MAP[iter_type]
-
-
-def get_itervar_label_color(itervar, iv_type):
- type_info = get_itervar_type_info(iv_type)
- return (
- linebrk(str(itervar["name"]) + "(" + type_info[0] + ")", TVMDD_TABLE_BODY_WIDTH),
- type_info[1],
- )
-
-
-def linebrk(s, n):
- """Break input string s with <br/> for every n charactors."""
- result = ""
- j = 0
- for i, c in enumerate(s):
- if j == n and i != len(s) - 1:
- result = result + "\n"
- j = 0
- j = j + 1
- result = result + c
- result = html.escape(str(result), quote=True)
- result = result.replace("\n", "<br/>")
- return result
-
-
-def create_graph(name="", rankdir="BT"):
- graph = Digraph(name=name)
- graph.graph_attr["rankdir"] = rankdir
- return graph
-
-
-def itervar_label(itervar, index, index_color, label):
- return (
- '<TR><TD PORT="'
- + itervar["id"]
- + '" BGCOLOR="'
- + index_color
- + '">'
- + str(index)
- + '</TD><TD BGCOLOR="white" PORT="itervar">'
- + label
- + "<br/>"
- + str(itervar["properties"]["range"])
- + "</TD></TR>"
- )
-
-
-def stage_label(stage):
- return stage["name"] + "<br/>Scope: " + stage["properties"]["scope"]
-
-
-def legend_label():
- """Generate legend labels."""
- label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">'
- for iter_type in ITERVAR_TYPE_STRING_MAP:
- name, color = ITERVAR_TYPE_STRING_MAP[iter_type]
- label += (
- '<TR><TD BGCOLOR="' + color + '"></TD>' + '<TD BGCOLOR="white">' + name + "</TD></TR>"
- )
- label += "</TABLE>>"
- return label
-
-
-def leaf_itervars(stage):
- filtered = filter(lambda x: (x["index"] >= 0), stage["all_itervars"])
- return sorted(filtered, key=lambda x: x["index"])
-
-
-def legend_dot(g):
- with g.subgraph(name="cluster_legend") as subgraph:
- subgraph.attr(label="Legend")
- label = legend_label()
- subgraph.node("legend", label, shape="none", margin="0")
-
-
-def extract_dom_for_viz(sch, need_range=True):
- json_str = dump_json(sch, need_range)
- s = json.loads(json_str)
- s = insert_dot_id(s)
- return s
-
-
-def dump_graph(dot_string, show_svg=True, dot_file_path="", output_dot_string=False):
- """Output dot_string in various formats."""
- if dot_file_path:
- try:
- dot_file = open(dot_file_path, "w+")
- dot_file.write(dot_string)
- dot_file.close()
- except IOError:
- print("Cannot open file: " + dot_file_path)
- if show_svg:
- from IPython.display import display
- from IPython.display import SVG
-
- src = Source(dot_string)
- display(SVG(src.pipe(format="svg")))
- if output_dot_string:
- return dot_string
- return None
-
-
-def dump_json(sch, need_range):
- """Serialize data for visualization from a schedule in JSON format.
-
- Parameters
- ----------
- sch : schedule
- The schedule object to serialize
-
- Returns
- -------
- json : string
- Serialized JSON string
- """
-
- def encode_itervar(itervar, stage, index, range_map):
- """Extract and encode IterVar visualization data to a dictionary"""
- ivrange = range_map[itervar] if range_map is not None and itervar in range_map else None
- bind_thread = None
- tensor_intrin = None
- if itervar in stage.iter_var_attrs:
- attr = stage.iter_var_attrs[itervar]
- iv_type = attr.iter_type
- # binding
- bind_thread = str(attr.bind_thread.var) if attr.bind_thread is not None else None
- # tensorization
- if attr.tensor_intrin is not None:
- tensor_intrin = str(attr.tensor_intrin.body)
- # remove the final \n
- tensor_intrin = tensor_intrin[0:-1] if tensor_intrin[-1] == "\n" else tensor_intrin
- else:
- tensor_intrin = None
- else:
- iv_type = itervar.iter_type
- itervar_dict = {
- "type": "IterVar",
- "index": index,
- "name": str(itervar.var),
- "itervar_type": iv_type,
- "properties": {
- "thread": bind_thread,
- "intrin": tensor_intrin,
- "range": str(ivrange) if ivrange is not None else "range(N/A)",
- },
- }
- return itervar_dict
-
- def encode_itervars(stage, range_map):
- """Extract and encode IterVars visualization data from a stage to a dictionary"""
-
- def get_leaf_itervar_index(itervar, leaf_iv):
- for leaf_index, ivar in enumerate(leaf_iv):
- if itervar_equal(ivar, itervar):
- return leaf_index
- return -1
-
- itervars = []
- for itervar in stage.all_iter_vars:
- leaf_index = get_leaf_itervar_index(itervar, stage.leaf_iter_vars)
- itervars.append(encode_itervar(itervar, stage, leaf_index, range_map))
- return itervars
-
- def encode_itervar_relation(obj_manager, rel):
- """Extract and encode IterVar Relationship visualization data to a dictionary"""
- rel_type = type(rel)
- if rel_type is tvm.te.schedule.Split:
- node_type = "Split_Relation"
- rel_dict = {
- "type": node_type,
- "parent": obj_manager.get_dom_path(rel.parent),
- "outer": obj_manager.get_dom_path(rel.outer),
- "inner": obj_manager.get_dom_path(rel.inner),
- }
- elif rel_type is tvm.te.schedule.Fuse:
- node_type = "Fuse_Relation"
- rel_dict = {
- "type": node_type,
- "fused": obj_manager.get_dom_path(rel.fused),
- "outer": obj_manager.get_dom_path(rel.outer),
- "inner": obj_manager.get_dom_path(rel.inner),
- }
- elif rel_type is tvm.te.schedule.Singleton:
- node_type = "Singleton_Relation"
- rel_dict = {
- "type": node_type,
- "iter": obj_manager.get_dom_path(rel.iter),
- }
- else:
- return None
- return rel_dict
-
- def encode_itervar_relations(obj_manager, stage):
- relations = []
- for i in range(len(stage.relations)):
- rel = encode_itervar_relation(obj_manager, stage.relations[i])
- if rel is not None:
- relations.append(rel)
- return relations
-
- def encode_tensor(obj_manager, tensor, stage):
- """Extract and encode tensor visualization data to a dictionary"""
- tensor_dict = {
- "type": "Tensor",
- "source": obj_manager.get_dom_path(stage),
- "value_index": tensor.value_index,
- "shape": str(tensor.op.output(tensor.value_index).shape),
- "data_type": tensor.op.output(tensor.value_index).dtype,
- }
- return tensor_dict
-
- def encode_tensors(obj_manager, stage):
- tensors = []
- for i in range(stage.op.num_outputs):
- tensor = stage.op.output(i)
- tensors.append(encode_tensor(obj_manager, tensor, stage))
- tensors.sort(key=lambda tensor: tensor["value_index"])
- return tensors
-
- def encode_stage(obj_manager, stage, range_map):
- """Extract and encode stage visualization data to a dictionary"""
- stage_dict = {
- "type": "Stage",
- "name": stage.op.name,
- "attaching_to": obj_manager.get_dom_path(stage.attach_ivar),
- "compute": str(stage.op.body) if hasattr(stage.op, "body") else None,
- "properties": {
- "scope": stage.scope,
- },
- "all_itervars": encode_itervars(stage, range_map),
- "relations": encode_itervar_relations(obj_manager, stage),
- "input_tensors": [
- obj_manager.get_dom_path(frozenset({tensor.op.name, tensor.value_index}))
- for tensor in stage.op.input_tensors
- ],
- "output_tensors": encode_tensors(obj_manager, stage),
- }
- return stage_dict
-
- def encode_schedule(sch, need_range):
- """Extract and encode data from a schedule for visualization to a nested dictionary.
- It is useful for JSON to serialize schedule.
-
- Parameters
- ----------
- sch : schedule
- The schedule object to extract
-
- Returns
- -------
- dict : dictionary
- A nested dictionary
- """
- assert isinstance(
- sch, tvm.te.schedule.Schedule
- ), "Input is not a tvm.te.schedule.Schedule object."
- range_map = None
- if need_range:
- try:
- range_map = tvm.te.schedule.InferBound(sch)
- except tvm._ffi.base.TVMError as expt:
- warnings.warn(
- "Ranges are not available, because InferBound fails with the following error:\n"
- + str(expt)
- )
-
- obj_manager = ObjectManager(sch)
- stages = []
- for stage in sch.stages:
- stages.append(encode_stage(obj_manager, stage, range_map))
- return {
- "type": "Schedule",
- "stages": stages,
- }
-
- return json.dumps(sch, default=lambda s: encode_schedule(s, need_range))
-
-
-def viz_schedule_tree(sch, show_svg=False, dot_file_path="", output_dot_string=False):
- """Top level API to render schedule tree
-
- Parameters
- ----------
- sch : schedule
- The schedule object to visualize
-
- show_svg : bool
- Display graph as SVG, useful for Jupyter notebooks.
-
- dot_file_path : string
- Dot file to save the graph.
-
- output_dot_string : bool
- Return dot file content or an empty string.
-
- Returns
- -------
- dot_string : string
- Dot file content or an empty string according to output_dot_string
-
- Examples
- --------
- The following code writes a schedule tree to a dot file.
-
- .. code-block:: python
- tedd.viz_schedule_tree(s, dot_file_path = '/tmp/example.dot')
-
- Use the following code to render a SVG graph in a Jupyter notebook.
-
- .. code-block:: python
- tedd.viz_schedule_tree(s, show_svg = True)
- """
-
- def create_schedule_tree_graph(name=""):
- return create_graph(name=name, rankdir="BT")
-
- def root_dot(g):
- g.node("ROOT", "ROOT", shape="oval", margin="0")
-
- def stage_node_dot(g, stage):
- node_label = stage_node_label(stage)
- g.node(stage["id"], node_label, shape="none", margin="0")
-
- def stage_node_label(stage):
- """Return a html format label for the given stage."""
- label = (
- '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" '
- 'CELLPADDING="4"> <TR><TD BGCOLOR="lightgrey" '
- 'COLSPAN="2" PORT="stage">' + stage_label(stage) + "</TD></TR>"
- )
-
- for leafiv in leaf_itervars(stage):
- iv_type = leafiv["itervar_type"]
- var_attr_label = ""
- if "thread" in leafiv["properties"] and leafiv["properties"]["thread"] is not None:
- var_attr_label = (
- var_attr_label
- + '<br/><font color="#2980B9">('
- + str(leafiv["properties"]["thread"])
- + ")</font>"
- )
- if "intrin" in leafiv["properties"] and leafiv["properties"]["intrin"] is not None:
- var_attr_label = (
- var_attr_label
- + "<br/>"
- + linebrk(
- "(tensor_intrin:" + str(leafiv["properties"]["intrin"]) + ")",
- TVMDD_TABLE_BODY_WIDTH,
- )
- )
- var_label, color = get_itervar_label_color(leafiv, iv_type)
- label += itervar_label(leafiv, leafiv["index"], color, var_label + var_attr_label)
- if stage["compute"] is not None:
- label += (
- '<TR><TD COLSPAN="2">'
- + linebrk(str(stage["compute"]), TVMDD_TABLE_BODY_WIDTH)
- + "</TD></TR>"
- )
- label += "</TABLE>>"
- return label
-
- def compute_at_dot(g, stage):
- """If the given stage attaches to another stage, create an edge from it
- stage to its attach point; otherwise, create an edge to the ROOT.
- """
- src = stage["id"]
- dst = (
- dom_path_to_string([stage["attaching_to"][0]], "Stage")
- + ":"
- + dom_path_to_string(stage["attaching_to"], "IterVar")
- if stage["attaching_to"] is not None
- else "ROOT"
- )
- color = (
- PALETTE[stage["attaching_to"][1] + 1]
- if stage["attaching_to"] is not None and stage["attaching_to"][1] < PALETTE_SIZE - 1
- else PALETTE[0]
- )
- g.edge(src, dst, color=color)
-
- graph = create_schedule_tree_graph("Schedule Tree")
- s = extract_dom_for_viz(sch)
- legend_dot(graph)
- for stage in s["stages"]:
- stage_node_dot(graph, stage)
- for stage in s["stages"]:
- compute_at_dot(graph, stage)
- root_dot(graph)
- return dump_graph(graph.source, show_svg, dot_file_path, output_dot_string)
-
-
-def viz_itervar_relationship_graph(sch, show_svg=False, dot_file_path="", output_dot_string=False):
- """Top level API to render IterVar relationship graph
-
- Parameters
- ----------
- sch : schedule
- The schedule object to visualize
-
- show_svg : bool
- Display graph as SVG, useful for Jupyter notebooks.
-
- dot_file_path : string
- Dot file to save the graph.
-
- output_dot_string : bool
- Return dot file content or an empty string.
-
- Examples
- --------
- The following code writes Ian tervar relationship graph to a dot file.
-
- .. code-block:: python
- tedd.viz_def viz_itervar_relationship_graph(sch,
- (s, dot_file_path = '/tmp/example.dot')
-
- Use the following code to render a SVG graph in a Jupyter notebook.
-
- .. code-block:: python
- tedd.viz_def viz_itervar_relationship_graph(sch,
- (s, show_svg = True)
- """
-
- def create_itervar_relation_graph(name=""):
- return create_graph(name=name, rankdir="TB")
-
- def itervar_node_dot(g, itervar, iv_type, index):
- label = itervar_node_label(itervar, iv_type, index)
- g.node(itervar["id"], label, shape="none", margin="0")
-
- def itervar_node_label(itervar, iv_type, index):
- label = (
- '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" '
- 'CELLPADDING="4">'
- + itervar_label(
- itervar,
- index,
- get_itervar_label_color(itervar, iv_type)[1],
- get_itervar_label_color(itervar, iv_type)[0],
- )
- + "</TABLE>>"
- )
- return label
-
- def itervar_relation_node_dot(g, node_id, node_label, input_ports, output_ports):
- label = itervar_relation_node_label(node_label, input_ports, output_ports)
- g.node(node_id, label, shape="none", margin="0")
-
- def itervar_relation_node_label(node_label, input_ports, output_ports):
- """Return a html format label for an itervar relationship node
- including node_label and input/output ports.
- """
- label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" ' 'CELLPADDING="4">' + "<TR>"
- max_port_num = max(len(input_ports), len(output_ports))
- for i in range(max_port_num):
- if i < len(input_ports):
- input_port = input_ports[i]
- label += '<TD BGCOLOR="lightgrey" PORT="' + input_port + '">' + input_port + "</TD>"
- else:
- label += '<TD BGCOLOR="white"></TD>'
- label += "</TR>"
- label += (
- '<TR><TD BGCOLOR="white" COLSPAN="'
- + str(max_port_num)
- + '" PORT="relation">'
- + node_label
- + "</TD></TR>"
- )
- label += "<TR>"
- for i in range(max_port_num):
- if i < len(output_ports):
- output_port = output_ports[i]
- label += (
- '<TD BGCOLOR="lightgrey" PORT="' + output_port + '">' + output_port + "</TD>"
- )
- else:
- label += '<TD BGCOLOR="white"></TD>'
- label += "</TR>"
- label += "</TABLE>>"
- return label
-
- def itervar_relation_dot(g, node, node_id):
- """Create an itervar relationship node."""
- node_type = node["type"]
- if node_type == "Split_Relation":
- node_type = "Split"
- itervar_relation_node_dot(g, node_id, node_type, ["Input"], ["Outer", "Inner"])
- parent = dom_path_to_string(node["parent"], "IterVar")
- outer = dom_path_to_string(node["outer"], "IterVar")
- inner = dom_path_to_string(node["inner"], "IterVar")
- g.edge(parent + ":itervar", node_id + ":Input")
- g.edge(node_id + ":Outer", outer + ":itervar")
- g.edge(node_id + ":Inner", inner + ":itervar")
- elif node_type == "Fuse_Relation":
- node_type = "Fuse"
- itervar_relation_node_dot(g, node_id, node_type, ["Outer", "Inner"], ["Fused"])
- fused = dom_path_to_string(node["fused"], "IterVar")
- outer = dom_path_to_string(node["outer"], "IterVar")
- inner = dom_path_to_string(node["inner"], "IterVar")
- g.edge(outer + ":itervar", node_id + ":Outer")
- g.edge(inner + ":itervar", node_id + ":Inner")
- g.edge(node_id + ":Fused", fused + ":itervar")
- elif node_type == "Singleton_Relation":
- node_type = "Singleton"
- itervar_relation_node_dot(g, node_id, node_type, [], ["Iter"])
- itervar = dom_path_to_string(node["inner"], "IterVar")
- g.edge(node_id + ":Iter", itervar + ":itervar")
- else:
- assert False, "Unknown IterVarRelationNode: " + node_type
-
- def stage_node_dot(g, stage):
- """Create a stage node."""
- with g.subgraph(name="cluster_" + stage["id"]) as subgraph:
- subgraph.attr(label=stage["name"])
- if stage["all_itervars"]:
- for itervar in stage["all_itervars"]:
- iv_type = itervar["itervar_type"]
- itervar_node_dot(subgraph, itervar, iv_type, itervar["index"])
- for rel in stage["relations"]:
- node_id = rel["id"]
- itervar_relation_dot(subgraph, rel, node_id)
- else:
- subgraph.node(stage["name"] + "_placeholder", style="invis")
-
- graph = create_itervar_relation_graph("IterVar Relationship Graph")
- s = extract_dom_for_viz(sch)
- legend_dot(graph)
- for stage in s["stages"]:
- stage_node_dot(graph, stage)
-
- return dump_graph(graph.source, show_svg, dot_file_path, output_dot_string)
-
-
-def viz_dataflow_graph(sch, show_svg=False, dot_file_path="", output_dot_string=False):
- """Top level API to render dataflow graph
-
- Parameters
- ----------
- sch : schedule
- The schedule object to visualize
-
- show_svg : bool
- Display graph as SVG, useful for Jupyter notebooks.
-
- dot_file_path : string
- Dot file to save the graph.
-
- output_dot_string : bool
- Return dot file content or an empty string.
-
- Examples
- --------
- The following code writes a dataflow graph to a dot file.
-
- .. code-block:: python
- tedd.viz_dataflow_graph(s, dot_file_path = '/tmp/example.dot')
-
- Use the following code to render a SVG graph in a Jupyter notebook.
-
- .. code-block:: python
- tedd.viz_dataflow_graph(s, show_svg = True)"""
-
- def create_dataflow_graph(name=""):
- return create_graph(name=name, rankdir="LR")
-
- def tensor_node_dot(g, tensor):
- """Create a tensor node."""
- label = tensor_node_label(tensor)
- g.node(tensor["id"], label, shape="oval", margin="0")
-
- def tensor_node_label(tensor):
- """Return a html format label for the given tensor."""
- label = str(tensor["shape"]) + "\n" + str(tensor["data_type"])
- return label
-
- def stage_node_dot(g, stage):
- """Create a stage node."""
- label = stage_node_label(stage)
- g.node(stage["id"], label, shape="none", margin="0")
-
- def stage_node_label(stage):
- """Return a html format label for the given stage."""
- rows = max(1, max(len(stage["output_tensors"]), len(stage["input_tensors"])))
- label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" ' 'CELLPADDING="4">'
- for i in range(rows):
- label += "<TR>"
- if i < len(stage["input_tensors"]):
- port_id = get_port_id(True, i)
- label += (
- '<TD BGCOLOR="lightgrey" COLSPAN="2" PORT="' + port_id + '">' + str(i) + "</TD>"
- )
- else:
- label += '<TD BGCOLOR="white" COLSPAN="2"></TD>'
- if i == 0:
- label += (
- '<TD BGCOLOR="white" COLSPAN="2" ROWSPAN="'
- + str(rows)
- + '">'
- + stage_label(stage)
- + "</TD>"
- )
- if i < len(stage["output_tensors"]):
- port_id = get_port_id(False, i)
- label += (
- '<TD BGCOLOR="lightgrey" COLSPAN="2" PORT="' + port_id + '">' + str(i) + "</TD>"
- )
- else:
- label += '<TD BGCOLOR="white" COLSPAN="2"></TD>'
- label += "</TR>"
- label += "</TABLE>>"
- return label
-
- def dfg_dot(g, sch):
- """Create edges among stages."""
- stages = sch["stages"]
- for stage in stages:
- for i in range(len(stage["input_tensors"])):
- src = dom_path_to_string(stage["input_tensors"][i], "Tensor")
- dst = stage["id"] + ":" + get_port_id(True, i)
- g.edge(src, dst)
- for i in range(len(stage["output_tensors"])):
- src = stage["id"] + ":" + get_port_id(False, i)
- dst = stage["output_tensors"][i]["id"]
- g.edge(src, dst)
-
- graph = create_dataflow_graph("Dataflow Graph")
- s = extract_dom_for_viz(sch, need_range=False)
- for stage in s["stages"]:
- stage_node_dot(graph, stage)
- for tensor in stage["output_tensors"]:
- tensor_node_dot(graph, tensor)
-
- dfg_dot(graph, s)
-
- return dump_graph(graph.source, show_svg, dot_file_path, output_dot_string)
diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
index fb325de1d3..f0fbf73755 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -17,106 +17,37 @@
# pylint: disable=invalid-name
"""The build utils in python."""
-from typing import Union, Optional, List, Mapping
+from typing import Union, Optional
-import warnings
import tvm.tir
-from tvm import te
-from tvm.runtime import Module
from tvm.runtime import ndarray
from tvm.ir import container
from tvm.tir import PrimFunc
from tvm.ir.module import IRModule
-from tvm.te import tensor
from tvm.target import Target
-from tvm.tir.buffer import Buffer
-from tvm.tir.expr import Var
from tvm.driver import _ffi_api as _driver_ffi
from . import _ffi_api as ffi
-def get_binds(args, compact=False, binds=None):
- """Internal function to get binds and arg_list given arguments.
- Parameters
- ----------
- args : list of Buffer or Tensor or Var
- The argument lists to the function.
- compact : bool
- If the statement has already bound to a compact buffer.
- binds : dict of :any:`Tensor` to :any:`Buffer`, optional
- Dictionary that maps the Tensor to Buffer which specified the data layout
- requirement of the function. By default, a new compact buffer is created
- for each tensor in the argument.
- Returns
- -------
- binds: dict
- The bind specification
- arg_list: list
- The list of symbolic buffers of arguments.
- """
- binds, arg_list = ffi.get_binds(args, compact, binds)
- return binds, arg_list
-
-
-def schedule_to_module(
- sch: te.Schedule,
- args: Optional[List[Union[Buffer, tensor.Tensor, Var]]] = None,
- name: str = "main",
- binds: Optional[Mapping[tensor.Tensor, Buffer]] = None,
-) -> IRModule:
- """According to the given schedule, form a function.
-
- This is a low-level function intended for testing purposes, and
- does not apply any optimization passes. In general, `tvm.lower`
- and `tvm.build` should be used instead.
-
- Parameters
- ----------
- sch : tvm.te.schedule.Schedule
- The given scheduler to form the raw body
- args : list of Buffer or Tensor or Var
- The argument lists to the function.
- name : str
- The name of result function, default name is "main"
- binds : dict of :any:`Tensor` to :any:`Buffer`, optional
- The binds information
- Returns
- -------
- The body formed according to the given schedule
- """
- return ffi.schedule_to_module(sch, args, name, binds)
-
-
def lower(
- inp: Union[te.Schedule, PrimFunc, IRModule],
- args: Optional[List[Union[Buffer, tensor.Tensor, Var]]] = None,
+ inp: Union[PrimFunc, IRModule],
name: str = "main",
- binds: Optional[Mapping[tensor.Tensor, Buffer]] = None,
simple_mode: bool = False,
) -> IRModule:
"""Lowering step before build into target.
Parameters
----------
- inp : Union[tvm.te.schedule.Schedule, tvm.tir.PrimFunc, IRModule]
+ inp : Union[tvm.tir.PrimFunc, IRModule]
The TE schedule or TensorIR PrimFunc/IRModule to be built
- args : Optional[List[Union[tvm.tir.Buffer, tensor.Tensor, tir.Var]]]
- The argument lists to the function for TE schedule.
-
- It should be None if we want to lower TensorIR.
name : str
The name of the result function.
- binds : Optional[Mapping[tensor.Tensor, tvm.tir.Buffer]]
- Dictionary that maps the Tensor to Buffer which specified the data layout
- requirement of the function. By default, a new compact buffer is created
- for each tensor in the argument.
-
simple_mode : bool
Whether only output simple and compact statement, this will skip
LoopPartition, api wrapper generation and Unrolling.
@@ -130,107 +61,40 @@ def lower(
return ffi.lower_module(inp, simple_mode)
if isinstance(inp, PrimFunc):
return ffi.lower_primfunc(inp, name, simple_mode)
- if isinstance(inp, te.Schedule):
- return ffi.lower_schedule(inp, args, name, binds, simple_mode)
raise ValueError(
f"Expected input to be an IRModule, PrimFunc or te.Schedule, but got {type(inp)}"
)
def build(
- inputs: Union[te.Schedule, PrimFunc, IRModule, Mapping[str, IRModule]],
- args: Optional[List[Union[Buffer, tensor.Tensor, Var]]] = None,
+ inputs: Union[PrimFunc, IRModule],
target: Optional[Union[str, Target]] = None,
- target_host: Optional[Union[str, Target]] = None,
- runtime: Optional[
- "tvm.relay.backend.Runtime"
- ] = None, # Type is annotated this way to avoid cyclic dependency
- name: Optional[str] = "default_function",
- binds: Optional[Mapping[tensor.Tensor, Buffer]] = None,
+ name: str = "main",
):
"""Build a function with arguments as signature. Code will be generated
for devices coupled with target information.
Parameters
----------
- inputs : Union[tvm.te.schedule.Schedule, tvm.tir.PrimFunc, IRModule, Mapping[str, IRModule]]
+ input : Union[tvm.tir.PrimFunc, IRModule]
The input to be built
- args : Optional[List[Union[tvm.tir.Buffer, tensor.Tensor, tir.Var]]]
- The argument lists to the function.
-
target : Optional[Union[str, Target]]
The target and option of the compilation.
- target_host : Optional[Union[str, Target]]
- Host compilation target, if target is device.
- When TVM compiles device specific program such as CUDA,
- we also need host(CPU) side code to interact with the driver
- setup the dimensions and parameters correctly.
- target_host is used to specify the host side codegen target.
- By default, llvm is used if it is enabled,
- otherwise a stackvm interpreter is used.
-
- runtime : Optional[Runtime]
- Runtime to generate artifacts for
-
- name : Optional[str]
+ name : str
The name of result function.
- binds : Optional[Mapping[tensor.Tensor, tvm.tir.Buffer]]
- Dictionary that maps the binding of symbolic buffer to Tensor.
- By default, a new buffer is created for each tensor in the argument.
-
Returns
-------
ret : tvm.module
A module that combines both host and device code.
- Examples
- ________
- There are two typical example uses of this function depending on the type
- of the argument `inputs`:
- 1. it is an IRModule.
-
- .. code-block:: python
-
- n = 2
- A = te.placeholder((n,), name='A')
- B = te.placeholder((n,), name='B')
- C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
- s = tvm.te.create_schedule(C.op)
- m = tvm.lower(s, [A, B, C], name="test_add")
- rt_mod = tvm.build(m, target="llvm")
-
- 2. it is a dict of compilation target to IRModule.
-
- .. code-block:: python
-
- n = 2
- A = te.placeholder((n,), name='A')
- B = te.placeholder((n,), name='B')
- C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
- s1 = tvm.te.create_schedule(C.op)
- with tvm.target.cuda() as cuda_tgt:
- s2 = topi.cuda.schedule_injective(cuda_tgt, [C])
- m1 = tvm.lower(s1, [A, B, C], name="test_add1")
- m2 = tvm.lower(s2, [A, B, C], name="test_add2")
- rt_mod = tvm.build({"llvm": m1, "cuda": m2})
-
Note
----
See the note on :any:`tvm.target` on target string format.
"""
- if isinstance(inputs, te.Schedule):
- if args is None:
- raise ValueError("args must be given for build from schedule")
- input_mod = lower(inputs, args, name=name, binds=binds)
- elif isinstance(inputs, (list, tuple, container.Array)):
- merged_mod = tvm.IRModule({})
- for x in inputs:
- merged_mod.update(lower(x))
- input_mod = merged_mod
- elif isinstance(inputs, PrimFunc):
+ if isinstance(inputs, PrimFunc):
input_mod = lower(inputs, name=name)
elif isinstance(inputs, tvm.IRModule):
assert (
@@ -238,31 +102,24 @@ def build(
), "Expected a non-empty IRModule, but the IRModule contained no functions."
input_mod = lower(inputs)
elif not isinstance(inputs, (dict, container.Map)):
- raise ValueError(
- f"Inputs must be te.Schedule, IRModule, PrimFunc, "
- f"or dict of target to IRModule, "
- f"but got {type(inputs)}."
- )
-
- if not isinstance(inputs, (dict, container.Map)):
- target = Target.current() if target is None else target
- if target is None and isinstance(input_mod, tvm.IRModule):
- target_mod = {}
- for gvar, func in input_mod.functions.items():
- tgt = func.attrs["target"] if "target" in func.attrs else "llvm"
- if tgt not in target_mod:
- target_mod[tgt] = {}
- target_mod[tgt][gvar] = func
-
- target_input_mod = {}
- for tgt in target_mod.keys():
- tir_mod = tvm.IRModule(target_mod[tgt])
- tir_mod = tir_mod.with_attrs(input_mod.attrs)
- target_input_mod[tgt] = tir_mod
- else:
- target_input_mod = {target: input_mod}
+ raise ValueError("Inputs must be IRModule or PrimFunc")
+
+ target = Target.current() if target is None else target
+ if target is None and isinstance(input_mod, tvm.IRModule):
+ target_mod = {}
+ for gvar, func in input_mod.functions.items():
+ tgt = func.attrs["target"] if "target" in func.attrs else "llvm"
+ if tgt not in target_mod:
+ target_mod[tgt] = {}
+ target_mod[tgt][gvar] = func
+
+ target_input_mod = {}
+ for tgt in target_mod.keys():
+ tir_mod = tvm.IRModule(target_mod[tgt])
+ tir_mod = tir_mod.with_attrs(input_mod.attrs)
+ target_input_mod[tgt] = tir_mod
else:
- target_input_mod = {tgt: lower(mod) for tgt, mod in inputs.items()}
+ target_input_mod = {target: input_mod}
# Because modules can be created from a variety of sources, we annotate them
# with the relevant attributes here to ensure they propagate
@@ -271,18 +128,10 @@ def build(
if not isinstance(tgt, (str, Target)):
raise ValueError("The key of inputs must be str or " "Target when inputs is dict.")
if not isinstance(mod, tvm.IRModule):
- raise ValueError("inputs must be Schedule, IRModule, " "or dict of str to IRModule.")
- annotated_mods[tgt] = mod.with_attr("runtime", runtime)
+ raise ValueError("inputs must be IRModule, " "or dict of str to IRModule.")
+ annotated_mods[tgt] = mod
- # TODO(mbs): Both CompilationConfig and TIRToRuntime implement the same host target
- # defaulting logic, but there's currently no way to get back the decided host.
- if target_host is not None:
- warnings.warn(
- "target_host parameter is going to be deprecated. "
- "Please pass in tvm.target.Target(target, host=target_host) instead."
- )
-
- annotated_mods, target_host = Target.canon_target_map_and_host(annotated_mods, target_host)
+ annotated_mods, target_host = Target.canon_target_map_and_host(annotated_mods)
if not target_host:
for tar, mod in annotated_mods.items():
device_type = ndarray.device(tar.kind.name, 0).device_type
@@ -296,41 +145,4 @@ def build(
rt_mod_host = _driver_ffi.tir_to_runtime(annotated_mods, target_host)
- annotated_mods, target_host = Target.canon_target_map_and_host(annotated_mods, target_host)
-
- if not isinstance(target_host, Target):
- target_host = Target(target_host)
-
- if str(runtime) == "crt" and runtime["system-lib"]:
- if target_host.kind.name == "c":
- create_csource_crt_metadata_module = tvm._ffi.get_global_func(
- "runtime.CreateCSourceCrtMetadataModule"
- )
- to_return = create_csource_crt_metadata_module([rt_mod_host], target_host, runtime)
- elif target_host.kind.name == "llvm":
- create_llvm_crt_metadata_module = tvm._ffi.get_global_func(
- "runtime.CreateLLVMCrtMetadataModule"
- )
- to_return = create_llvm_crt_metadata_module([rt_mod_host], target_host, runtime)
- else:
- to_return = rt_mod_host
-
- return OperatorModule.from_module(to_return, ir_module_by_target=annotated_mods, name=name)
-
-
-class OperatorModule(Module):
- """Wraps the Module returned by tvm.build() and captures additional outputs of that function."""
-
- @classmethod
- def from_module(cls, mod, **kwargs):
- # NOTE(areusch): It is generally unsafe to continue using `mod` from this point forward.
- # If an exception occurs in cls.__init__, handle will be deleted. For this reason,
- # set mod.handle to None.
- handle = mod.handle
- mod.handle = None
- return cls(handle, **kwargs)
-
- def __init__(self, handle, ir_module_by_target=None, name=None):
- super(OperatorModule, self).__init__(handle)
- self.ir_module_by_target = ir_module_by_target
- self.name = name
+ return rt_mod_host
diff --git a/python/tvm/exec/measure_peak.py b/python/tvm/exec/measure_peak.py
deleted file mode 100644
index 178e600892..0000000000
--- a/python/tvm/exec/measure_peak.py
+++ /dev/null
@@ -1,52 +0,0 @@
-# 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.
-"""measure bandwidth and compute peak
-
-e.g.
-python3 -m tvm.exec.measure_peak --target cuda --rpc-host 127.0.0.1 --rpc-port 9090
-python3 -m tvm.exec.measure_peak --target opencl --target-host "llvm -mtriple=aarch64-linux-gnu" \
- --rpc-host $TVM_OPENCL_DEVICE_HOST --rpc-port 9090
-"""
-
-import argparse
-import logging
-
-from tvm.target import Target
-from ..contrib.peak import measure_peak_all
-
-
-def main():
- """Main function"""
- parser = argparse.ArgumentParser()
- parser.add_argument("--target", type=str, default="llvm", help="The build target")
- parser.add_argument(
- "--target-host", type=str, default=None, help="The host code compilation target"
- )
- parser.add_argument(
- "--rpc-host", type=str, default="127.0.0.1", help="the hostname of the server"
- )
- parser.add_argument("--rpc-port", type=int, default=9090, help="The port of the RPC")
-
- args = parser.parse_args()
- logging.basicConfig(level=logging.INFO)
-
- args.target, args.target_host = Target.canon_target_and_host(args.target, args.target_host)
- measure_peak_all(args.target, args.target_host, args.rpc_host, args.rpc_port)
-
-
-if __name__ == "__main__":
- main()
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index cfa4143b66..9ea08bfd47 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -179,10 +179,12 @@ def _vmcodegen(
raise ValueError(f"Unknown exec_mode {exec_mode}")
-def _autodetect_system_lib_req(
- target: Optional[tvm.target.Target] = None, system_lib: Optional[bool] = None
+def _auto_attach_system_lib_prefix(
+ tir_mod: tvm.IRModule,
+ target: Optional[tvm.target.Target] = None,
+ system_lib: Optional[bool] = None,
):
- """Automatically detect system lib requirement"""
+ """Automatically detect system lib req and attach prefix attr"""
if target is not None:
host = target if target.host is None else target.host
if system_lib is None:
@@ -191,9 +193,9 @@ def _autodetect_system_lib_req(
system_lib = True
if system_lib:
- # use packed-func to avoid relay dep.
- return tvm.get_global_func("relay.backend.CreateRuntime")("cpp", {"system-lib": system_lib})
- return None
+ if tir_mod.get_attr("system_lib_prefix") is None:
+ return tir_mod.with_attr("system_lib_prefix", "")
+ return tir_mod
def _vmlink(
@@ -246,11 +248,7 @@ def _vmlink(
relax_ext_libs = []
tir_ext_libs = []
if tir_mod is not None and len(tir_mod.get_global_vars()) > 0:
- lib = tvm.build(
- tir_mod,
- target=target,
- runtime=_autodetect_system_lib_req(target, system_lib),
- )
+ lib = tvm.build(_auto_attach_system_lib_prefix(tir_mod, system_lib), target=target)
for ext_mod in ext_libs:
if ext_mod.is_device_module:
tir_ext_libs.append(ext_mod)
diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py
index 0907ea2ebf..9ad4a5b57b 100644
--- a/python/tvm/te/__init__.py
+++ b/python/tvm/te/__init__.py
@@ -28,21 +28,11 @@ from tvm.tir import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floor
from tvm.tir import comm_reducer, min, max, sum
from tvm.tir import add, subtract, multiply
-from .schedule import (
- Schedule,
- Stage,
- create_schedule,
- SpecializedCondition,
- AXIS_SEPARATOR,
-)
from .tensor import TensorSlice, Tensor
-from .tensor_intrin import decl_tensor_intrin
from .tag import tag_scope
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc
-from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
-from .autodiff import gradient
-from . import hybrid
+from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp
diff --git a/python/tvm/te/autodiff.py b/python/tvm/te/autodiff.py
deleted file mode 100644
index f865083994..0000000000
--- a/python/tvm/te/autodiff.py
+++ /dev/null
@@ -1,67 +0,0 @@
-# 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.
-
-"""Automatic differentiation of tensor expressions."""
-from . import _ffi_api
-
-
-def gradient(output, inputs, head=None):
- """Perform reverse-mode automatic differentiation.
-
- Parameters
- ----------
- output : Tensor
- The tensor to differentiate.
-
- inputs : List[Tensor]
- The list of input tensors to be differentiated wrt.
-
- head : Tensor
- The adjoint of the output, in other words, some tensor, by which the Jacobians
- will be multiplied. Its shape must be of the form `prefix + output.shape`.
- If `None` is passed, the identity tensor of shape `output.shape + output.shape`
- will be used.
-
- Returns
- -------
- tensors: List[Tensor]
- The result gradient, in the same order as the inputs
-
- Example
- -------
- .. code-block:: python
-
- x = tvm.placeholder((32, 3, 28, 28), name='x')
- w1 = tvm.placeholder((10, 3, 3, 3), name='w1')
- w2 = tvm.placeholder((10, 10, 3, 3), name='w2')
- z1 = topi.nn.conv2d(x, w1, 1, 1, 1)
- z2 = topi.nn.conv2d(z1, w2, 1, 1, 1)
- y = topi.sum(z2)
-
- # produce gradients
- [dw1, dw2] = tvm.gradient(y, [w1, w2])
-
- # produce Jacobians
- [jw1, jw2] = tvm.gradient(z2, [w1, w2])
-
- # produce gradients, the head adjoint for z2 is provided manually
- [dw1, dw2] = tvm.gradient(z2, [w1, w2], topi.full_like(z2, 1.0))
-
- """
- if not isinstance(inputs, list):
- inputs = [inputs]
- return _ffi_api.Gradient(output, inputs, head)
diff --git a/python/tvm/te/hybrid/__init__.py b/python/tvm/te/hybrid/__init__.py
deleted file mode 100644
index cd320c6b20..0000000000
--- a/python/tvm/te/hybrid/__init__.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# 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.
-"""Hybrid Programming APIs of TVM Python Package.
-
-This package maps a subset of python to HalideIR so that:
-1. Users can write some preliminary versions of the computation patterns
-have not been supported yet and verify it across the real execution and
-python semantic emulation.
-2. So far, it is a text format dedicated to HalideIR Phase 0. Refer tvm.lower
-for more details. A larger ambition of this module is to support all levels of
-HalideIR.
-"""
-
-# TODO(@were): Make this module more complete.
-# 1. Support HalideIR dumping to Hybrid Script
-# 2. Support multi-level HalideIR
-import inspect
-import tvm._ffi
-import tvm.te.schedule
-from tvm._ffi.base import decorate
-
-from .module import HybridModule
-from .parser import source_to_op
-from .utils import _pruned_source
-
-
-def script(pyfunc):
- """Decorate a python function as hybrid script.
-
- The hybrid function support emulation mode and parsing to
- the internal language IR.
-
- Returns
- -------
- hybrid_func : function
- A decorated hybrid script function.
- """
- # pylint: disable=import-outside-toplevel, missing-docstring
- def wrapped_func(func, *args, **kwargs):
- from .utils import _is_tvm_arg_types
-
- if _is_tvm_arg_types(args):
- src = _pruned_source(func)
- closure_vars = inspect.getclosurevars(func).nonlocals
- closure_vars.update(inspect.getclosurevars(func).globals)
- return source_to_op(src, args, func.__globals__, closure_vars)
-
- from .runtime import _enter_hybrid_runtime, _restore_runtime
-
- intersect = _enter_hybrid_runtime(func)
- value = func(*args, **kwargs)
- _restore_runtime(func, intersect)
- return value
-
- return decorate(pyfunc, wrapped_func)
-
-
-def build(sch, inputs, outputs, name="hybrid_func"):
- """Dump the current schedule to hybrid module
-
- Parameters
- ----------
- sch: tvm.te.Schedule
- The schedule to be dumped
-
- inputs: An array of Tensors or Vars
- The inputs of the function body
-
- outputs: An array of Tensors
- The outputs of the function body
-
- Returns
- -------
- module: HybridModule
- The built results is wrapped in a HybridModule.
- The usage of HybridModule is roughly the same as normal TVM-built modules.
- """
- sch = sch.normalize()
- bounds = tvm.te.schedule.InferBound(sch)
- stmt = tvm.te.schedule.ScheduleOps(sch, bounds)
-
- src = _Dump(stmt, inputs, outputs, name)
-
- return HybridModule(src, name)
-
-
-tvm._ffi._init_api("tvm.hybrid", __name__)
diff --git a/python/tvm/te/hybrid/calls.py b/python/tvm/te/hybrid/calls.py
deleted file mode 100644
index 948a0d7665..0000000000
--- a/python/tvm/te/hybrid/calls.py
+++ /dev/null
@@ -1,183 +0,0 @@
-# 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.
-"""Intrinsics of TVM-Python Hybrid Script for Python compilation time
-semantic support."""
-
-from tvm.runtime import const, convert
-import tvm.te
-from tvm.ir.container import Array
-from tvm.target import Target
-from tvm.tir import expr as _expr
-from tvm.tir import call_intrin
-from tvm.tir.stmt import ForKind
-
-from .utils import _internal_assert
-
-# pylint: disable=redefined-builtin,invalid-name
-
-LOOP_INTRIN = {
- "range": ForKind.SERIAL,
- "unroll": ForKind.UNROLLED,
- "parallel": ForKind.PARALLEL,
- "vectorize": ForKind.VECTORIZED,
- "const_range": (ForKind.UNROLLED,),
-}
-
-
-def _range(annotation, args):
- """Handling TVM loop types"""
- n = args.__len__()
- if n == 1:
- low, ext = const(0, dtype="int32"), args[0]
- else:
- _internal_assert(n == 2, "A loop intrinsic should only have 1 or 2 arguments!")
- low, ext = args[0], args[1]
- if not tvm.tir.analysis.expr_deep_equal(low, const(0, dtype="int32")):
- ext = ext - low
- kind = LOOP_INTRIN[annotation]
- iter_var = None
- return iter_var, low, ext, kind
-
-
-range = unroll = vectorize = parallel = const_range = _range # pylint: disable=invalid-name
-
-
-def bind(func_id, args):
- """Handling TVM thread binding"""
- _internal_assert(func_id == "bind", "This function cannot be directly invoked!")
- _internal_assert(args.__len__() == 2, "A loop bind should only have 2 arguments!")
- _internal_assert(isinstance(args[0], str), "A loop bind's first argument should be a string!")
- low, ext = const(0, "int32"), args[1]
- iter_var = tvm.te.thread_axis((low, ext), args[0])
- kind = None
- return iter_var, low, ext, kind
-
-
-def _math_intrin(func_id, args):
- # pylint: disable=import-outside-toplevel
- from tvm.tir import op
-
- return getattr(op, func_id)(*args)
-
-
-sqrt = (
- log
-) = exp = tanh = sigmoid = power = popcount = round = _math_intrin # pylint: disable=invalid-name
-
-
-def _min_max(func_id, args):
- _internal_assert(args.__len__() == 2, "Max/Min function should have 2 elements")
- return getattr(_expr, func_id.title())(args[0], args[1])
-
-
-min = max = _min_max # pylint: disable=invalid-name
-
-
-def _allocate_tensor(func_id, args):
- """Handling TVM tensor allocation.
- You may refer hybrid.intrin.allocate for more details."""
- n = args.__len__()
- _internal_assert(
- isinstance(convert(args[0]), Array), "allocate's first argument should be a tuple of shape!"
- )
- shape = args[0]
- for i in shape:
- _internal_assert(isinstance(i, (_expr.PrimExpr, int)), "The shape should be an expression")
- if n > 1:
- _internal_assert(isinstance(args[1], str), "The data type should be an str")
- _internal_assert(
- args[1].startswith("int") or args[1].startswith("float"),
- "The data type should be either int or float!",
- )
- dtype = args[1]
- else:
- dtype = "float32"
- if n > 2:
- _internal_assert(isinstance(args[2], str), "The data scope should be an string")
- _internal_assert(func_id != "output_tensor", "Output tensor cannot specify scope")
- scope = args[2]
- else:
- scope = "global" if func_id != "output_tensor" else "output"
- return (shape, dtype, scope)
-
-
-output_tensor = allocate = _allocate_tensor # pylint: disable=invalid-name
-
-
-def len(func_id, args):
- """Iterpret the len function"""
- _internal_assert(args.__len__() == 1, "Only 1 argument is expected!")
- _internal_assert(func_id == "len", "This function cannot be directly invoked!")
- try:
- return convert(args[0].__len__())
- except: # pylint: disable=bare-except
- _internal_assert(args[0].shape.__len__() == 1, "Only one-dimension array can get len")
- return convert(args[0].shape[0])
-
-
-def _cast(func_id, args):
- _internal_assert(
- args.__len__() == 1,
- f"Casting to {func_id} only supports a single argument",
- )
- # The FFI can handle any conversion of `args[0]` into PrimExpr, if
- # required.
- return _expr.Cast(func_id, args[0])
-
-
-float16 = float32 = float64 = _cast # pylint: disable=invalid-name
-int8 = int16 = int32 = int64 = _cast # pylint: disable=invalid-name
-uint8 = uint16 = uint32 = uint64 = _cast # pylint: disable=invalid-name
-
-
-def ceil_div(func_id, args):
- _internal_assert(func_id == "ceil_div", "This function cannot be directly invoked!")
- _internal_assert(args.__len__() == 2, "2 arguments expected for division!")
- a, b = args
- return (a + b - 1) // b
-
-
-def likely(func_id, args):
- _internal_assert(args.__len__() == 1, "Only one expression can be likely")
- _internal_assert(func_id == "likely", "This function cannot be directly invoked!")
- return call_intrin(args[0].dtype, "tir.likely", *args)
-
-
-def max_num_threads(func_id, args):
- """Set the maximum number of threads."""
- _internal_assert(func_id == "max_num_threads", "This function cannot be directly invoked!")
- _internal_assert(args.__len__() <= 1, "At most one argument accepted!")
- if args.__len__() == 0:
- res = Target.current().max_num_threads
- else:
- _internal_assert(isinstance(args[0], _expr.IntImm), "In tvm bool should be uint")
- res = Target.current(args[0].value).max_num_threads
- return convert(res)
-
-
-def inf(func_id, args):
- """Infinity"""
- _internal_assert(func_id == "inf", "This function cannot be directly invoked!")
- _internal_assert(args.__len__() == 1, "One argument accepted!")
- return tvm.tir.max_value(args[0])
-
-
-def ninf(func_id, args):
- """Negative infinity"""
- _internal_assert(func_id == "ninf", "This function cannot be directly invoked!")
- _internal_assert(args.__len__() == 1, "One argument accepted!")
- return tvm.tir.min_value(args[0])
diff --git a/python/tvm/te/hybrid/module.py b/python/tvm/te/hybrid/module.py
deleted file mode 100644
index 729805b31b..0000000000
--- a/python/tvm/te/hybrid/module.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# 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.
-"""Methods and data structures to support dumping HalideIR to Hybrid Script.
-This allows users to do quick hack to generated HalideIR and cast it back to
-TVM modules.
-
-To enable this feature, you need to build with -DUSE_HYBRID_DUMP=ON.
-"""
-
-import ast
-
-from tvm.contrib import utils
-from .utils import _internal_assert
-from .utils import _is_tvm_arg_types
-from .parser import source_to_op
-
-
-class HybridModule(object):
- """The usage of Hybrid Module is very similar to conventional TVM module,
- but conventional TVM module requires a function body which is already fully
- lowered. This contradicts to the fact that Hybrid Module is originally a text
- format for Phase 0 HalideIR. Thus, a totally separated module is defined."""
-
- def __init__(self, src=None, name=None):
- """The constructor of this a hybrid module
-
- Parameters
- ----------
- src : str
- The source code of this module
-
- name : str
- The name of this module
- """
- self.src_ = self.name = self.func_ = self.root_ = None
- if src is not None:
- temp = utils.tempdir()
- dst = temp.relpath("script.py")
- with open(dst, "w") as f:
- f.write(f"import tvm\n@tvm.te.hybrid.script\n{src}")
-
- if name is not None:
- self.name = name
- self.load(dst)
-
- def __call__(self, *args):
- if _is_tvm_arg_types(args):
- return source_to_op(self.root_, args, globals(), {})
- return self.func_(*args)
-
- def get_source(self):
- return self.src_
-
- def save(self, path):
- if not path.endswith(".py"):
- path = path + ".py"
- with open(path, "w") as f:
- f.write(self.src_)
-
- def load(self, path):
- """Load the module from a python file
-
- Parameters
- ----------
- path : str
- Path to the given python file
- """
- with open(path, "r") as f:
- self.src_ = f.read()
-
- src = self.src_
-
- class FindFunc(ast.NodeVisitor):
- """Find the function in module to be loaded module."""
-
- # pylint: disable=invalid-name
- def __init__(self):
- self.name = None
- self.root = None
-
- def visit_FunctionDef(self, node):
- _internal_assert(self.name is None, "For now, only one function supported!")
- self.name = node.name
- _internal_assert(self.root is None, "For now, only one function supported!")
- self.root = node
-
- root = ast.parse(src)
- finder = FindFunc()
- finder.visit(root)
- _internal_assert(finder.name is not None and finder.root is not None, "No function found!")
- if self.name is None:
- self.name = finder.name
- self.root_ = finder.root
-
- _, local_ = {}, {}
- exec(self.src_, _, local_) # pylint: disable=exec-used
- local_.pop("tvm")
- assert len(local_) == 1
- self.func_ = list(local_.values())[0]
diff --git a/python/tvm/te/hybrid/parser.py b/python/tvm/te/hybrid/parser.py
deleted file mode 100644
index bd5a060cd0..0000000000
--- a/python/tvm/te/hybrid/parser.py
+++ /dev/null
@@ -1,658 +0,0 @@
-# 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.
-"""Hybrid Script Parser"""
-
-import ast
-import operator
-import logging
-import sys
-import numbers
-
-from enum import Enum
-from tvm.ir import Array, Range
-import tvm.runtime
-import tvm.tir
-import tvm.te
-import tvm.te._ffi_api
-import tvm.arith
-
-from tvm.tir import expr as _expr
-from tvm.tir import stmt as _stmt
-from tvm.te.tensor import Tensor, Operation
-from tvm.tir import all as _all
-from tvm.tir import any as _any
-
-from .utils import _internal_assert
-from . import calls
-from . import utils
-from .preprocessor import determine_variable_usage
-
-
-def concat_list_to_block(lst):
- """Concatenate a list of Python IR nodes to HalideIR Block"""
- if not lst:
- return utils.make_nop()
- n = len(lst)
- if n == 1:
- return lst[0]
- return _stmt.SeqStmt(lst)
-
-
-def visit_list_to_block(visit, lst):
- """Visit and concatenate a list of Python IR nodes to HalideIR Block"""
- lst = [visit(stmt) for stmt in lst if not utils.is_docstring(stmt)]
- lst = [stmt for stmt in lst if not tvm.ir.structural_equal(stmt, utils.make_nop())]
- if not lst:
- return utils.make_nop()
- return concat_list_to_block(lst)
-
-
-class Symbol(Enum):
- """Enumerates types in the symbol table"""
-
- Callable = 0
- Input = 1
- OutputBuffer = 2
- GlobalBuffer = 3
- LocalBuffer = 4
- SharedBuffer = 5
- ConstVar = 6
- BufferVar = 7
- LoopVar = 8
- ConstLoopVar = 9
- ThreadBind = 10
-
-
-def _floordiv(x, y):
- if isinstance(x, _expr.ExprOp) or isinstance(y, _expr.ExprOp):
- return tvm.tir.floordiv(x, y)
- return operator.floordiv(x, y)
-
-
-def _floormod(x, y):
- if isinstance(x, _expr.ExprOp) or isinstance(y, _expr.ExprOp):
- return tvm.tir.floormod(x, y)
- return operator.mod(x, y)
-
-
-class HybridParser(ast.NodeVisitor):
- """Python AST visitor pass which finally lowers it to HalideIR"""
-
- _binop_maker = {
- ast.Add: operator.add,
- ast.Sub: operator.sub,
- ast.Mult: operator.mul,
- ast.Div: operator.div if sys.version_info[0] == 2 else operator.truediv,
- ast.FloorDiv: _floordiv,
- ast.Mod: _floormod,
- ast.BitOr: operator.or_,
- ast.BitAnd: operator.and_,
- ast.BitXor: operator.xor,
- ast.Gt: operator.gt,
- ast.GtE: operator.ge,
- ast.Lt: operator.lt,
- ast.LtE: operator.le,
- ast.Eq: operator.eq,
- ast.NotEq: operator.ne,
- ast.And: _all,
- ast.Or: _any,
- }
-
- _unaryop_maker = {ast.USub: operator.neg, ast.Invert: operator.invert, ast.Not: operator.not_}
-
- def __init__(self, args, usage, symbols, closure_vars, func_name=None):
- """
- Parameters
- ----------
- args: A list of tvm.te.placeholder or te.var
- Provided by the user, the argument list of the function to be lowered.
-
- usage: A dict of variables used in last in this function
- Provided by last lower pass, which collects this information
-
- symbols : list of str
- The symbol list of the global context of the function.
-
- closure_vars: dict
- A dict of external name reference captured by this function.
-
- Returns
- -------
- func_name: str
- The name of the function to be lowered; if not provided,
- the compiler will use the name in the AST
- """
- self.args = list(args)
- self.usage = usage.copy()
-
- self.symbols = {} # Symbol table
- for k, v in symbols.items():
- if callable(v):
- self.add_symbol(k, Symbol.Callable, v)
-
- self.closure_vars = closure_vars
-
- self.binds = {} # Thread binds
- self.device = 0 # Is it generating device
-
- self.func_name = func_name # The name of the function to be lowered
- self.outputs = [] # Output tensors' name
- self.side_effect = set() # Tensors with side effects
- self.parsed_body = None # The parsed HalideIR body
- self.analyzer = tvm.arith.Analyzer()
- self.returned = False # If this function has a valid return
-
- def add_symbol(self, key, ty, val): # pylint: disable=invalid-name
- """Add value to the symbol table context"""
- if key in self.symbols.keys():
- old = str(self.symbols[key])
- new = str((ty, val))
- _internal_assert(False, f"Name conflict in symbol table! [{key}] {old} -> {new}")
-
- self.symbols[key] = ty, val
-
- if ty == Symbol.ThreadBind:
- if val.var.name not in self.binds.keys():
- self.binds[val.var.name] = val
- return
- val_ = self.binds[val.var.name]
- _internal_assert(
- tvm.tir.analysis.expr_deep_equal(val_.dom.extent, val.dom.extent),
- "Thread extents should be uniform!",
- )
- self.symbols[key] = ty, val_
-
- def wrap_up_realize(self, node, body):
- """Wrap up all the variables which will no longer be used"""
- to_pop = []
- for key, val in self.usage.items():
- _, level, _ = val
- if key not in self.symbols:
- # don't realize the symbols that are never visited
- continue
- if level != node:
- continue
- _internal_assert(key in self.symbols.keys(), f"Unknown symbol {key}!")
-
- ty, entry = self.symbols[key] # pylint: disable=invalid-name
- if ty in [Symbol.Input, Symbol.OutputBuffer]:
- continue
- if "Buffer" in ty.name:
- _buf = entry
- _scope = "global" if ty is Symbol.BufferVar else ty.name[:-6].lower()
- to_pop.append(key)
- else:
- continue
-
- if _scope == "global":
- body = self.wrap_up_binds(body)
-
- _domain = [Range.from_min_extent(0, i) for i in _buf.shape]
- _dtype = _buf.dtype
- _true = tvm.runtime.convert(True)
- body = tvm.tir.ProducerRealize(_buf, _domain, _true, body, tvm.runtime.convert(_scope))
-
- for elem in to_pop:
- self.symbols.pop(elem)
-
- return body
-
- def wrap_up_binds(self, body):
- for _, iter_var in self.binds.items():
- ext = iter_var.dom.extent
- body = tvm.tir.AttrStmt(iter_var, "thread_extent", ext, body)
- self.binds = {}
- return body
-
- # pylint: disable=invalid-name, missing-docstring
- def visit_Module(self, node):
- _internal_assert(
- len(node.body) == 1, "Only one-function source code will be fed to this parser!"
- )
- return self.visit(node.body[0])
-
- def visit_FunctionDef(self, node):
- _internal_assert(
- len(node.args.args) == len(self.args),
- "The number of arguments passed to the \
- function should be the same as it is defined!",
- )
- if self.func_name is None:
- self.func_name = node.name
- for idx, arg in enumerate(node.args.args):
- _attr = "id" if sys.version_info[0] < 3 else "arg" # To make py2 and 3 compatible
- self.add_symbol(getattr(arg, _attr), Symbol.Input, self.args[idx])
- res = visit_list_to_block(self.visit, node.body)
- res = self.wrap_up_realize(node, res)
- return self.wrap_up_binds(res)
-
- def visit_Expr(self, node):
- return self.visit(node.value)
-
- def visit_Name(self, node):
- name = node.id
- if sys.version_info[0] == 2 and name in ["True", "False"]:
- return tvm.runtime.convert(ast.literal_eval(name))
-
- if name in self.closure_vars:
- return tvm.runtime.convert(self.closure_vars[name])
-
- ty, entry = self.symbols[name]
- _internal_assert(name in self.symbols, f"Unknown symbol {name}!")
- if ty in [Symbol.LoopVar, Symbol.Input, Symbol.ConstLoopVar]:
- return entry
- if ty is Symbol.ThreadBind:
- return entry.var
- if ty is Symbol.ConstVar:
- return entry if isinstance(node.ctx, ast.Load) else None
- if ty is Symbol.BufferVar:
- if isinstance(node.ctx, ast.Load):
- return tvm.tir.ProducerLoad(entry, [tvm.runtime.const(0, "int32")])
- return entry, [tvm.runtime.const(0, "int32")]
- # Do I need any assertion here?
- return entry
-
- def visit_Num(self, node):
- if isinstance(node.n, numbers.Integral):
- dtype = "int32"
- elif isinstance(node.n, float):
- dtype = "float32"
- else:
- _internal_assert(
- isinstance(node.n, bool), "The data type should be one of (int, float, bool)"
- )
- dtype = "bool"
- return tvm.runtime.const(node.n, dtype)
-
- def visit_NameConstant(self, node):
- return tvm.tir.const(node.value)
-
- def visit_AugAssign(self, node):
- buf = self.visit(node.target)
- rhs = self.visit(node.value)
- if isinstance(buf, tuple):
- _internal_assert(len(buf) == 2, "LHS is supposed to be (buf, args)!")
- buf, args = buf
- else:
- args = [tvm.runtime.const(0, "int32")]
- _internal_assert(isinstance(buf, Tensor), "LHS is supposed to be Tensor!")
-
- read = tvm.tir.ProducerLoad(buf, args)
- value = HybridParser._binop_maker[type(node.op)](read, rhs)
-
- return tvm.tir.ProducerStore(buf, value, args)
-
- def visit_Assign(self, node):
- rhs = self.visit(node.value)
- if isinstance(rhs, Operation):
- rmap = {}
- _internal_assert(
- len(node.targets) == rhs.num_outputs, "Unable to detuple the outs to targets"
- )
- for i in range(rhs.num_outputs):
- _internal_assert(
- isinstance(node.targets[i], ast.Name),
- "You should bind a pure name to the tensors",
- )
- self.add_symbol(node.targets[i].id, Symbol.GlobalBuffer, rhs.output(i))
- rmap[rhs.outputs[i].op] = rhs.output(i)
- return utils.replace_io(rhs.body, rmap)
-
- _internal_assert(len(node.targets) == 1, "So far only one-valued assignment is supported!")
- lhs = node.targets[0]
- if isinstance(rhs, _expr.PrimExpr):
- rhs = self.analyzer.simplify(rhs)
- if isinstance(lhs, ast.Name):
- # TODO: support defined intermediate buffer later
- lhs_ = lhs
- lhs = lhs.id
- if lhs in self.symbols.keys():
- ty, _ = self.symbols[lhs]
- _internal_assert(ty != Symbol.LoopVar, "Loop variable cannot be overwritten!")
- decl, _, rw = self.usage[lhs]
- if decl == lhs_:
- _internal_assert(
- lhs not in self.symbols.keys(),
- "This value should not be defined before this point!",
- )
- if isinstance(rhs, tuple):
- shape, dtype, scope = rhs
- ph = tvm.te.placeholder(shape, dtype=dtype, name=lhs)
- self.add_symbol(lhs, getattr(Symbol, scope.title() + "Buffer"), ph)
- if scope == "output":
- self.outputs.append(lhs)
- return utils.make_nop()
- if isinstance(rhs, utils.halide_imm_types) and ast.Store not in rw:
- self.add_symbol(lhs, Symbol.ConstVar, rhs)
- else:
- _internal_assert(
- self.device == 0,
- "Single variable not supported in devices' side!\n"
- + "If you are using GPU, please allocate a 'local' spad "
- + "outside the bind body",
- )
- ph = tvm.te.placeholder((1,), dtype=rhs.dtype, name=lhs)
- self.add_symbol(lhs, Symbol.BufferVar, ph)
- lhs = self.visit(lhs_)
- if lhs is not None:
- buf, args = lhs
- return tvm.tir.ProducerStore(buf, rhs, args)
- return utils.make_nop()
-
- lhs, args = self.visit(lhs)
- _internal_assert(
- isinstance(lhs, Tensor), "An array access's LHS is expected to be a expr.Call!"
- )
- res = tvm.tir.ProducerStore(lhs, rhs, args)
- return res
-
- def visit_Index(self, node):
- if isinstance(node.value, ast.Tuple):
- return self.visit(node.value)
- return [self.visit(node.value)]
-
- def visit_Attribute(self, node):
- buf = self.visit(node.value)
- return getattr(buf, node.attr)
-
- def visit_Subscript(self, node):
- args = self.visit(node.slice)
- if sys.version_info >= (3, 9):
- if not isinstance(node.slice, ast.Tuple):
- args = [args]
-
- arr = self.visit(node.value)
- if isinstance(arr, (Array, list, tuple)):
- for i in args:
- if isinstance(i, numbers.Integral):
- arr = arr[i]
- else:
- _internal_assert(
- isinstance(i, (_expr.IntImm,)), "All indices are supposed to be constants"
- )
- arr = arr[i.value]
- return arr
- if isinstance(node.ctx, ast.Load):
- return tvm.tir.ProducerLoad(arr, args)
- return arr, args
-
- def visit_With(self, node):
- if sys.version_info[0] < 3:
- context = node.context_expr
- option = node.optional_vars
- else:
- _internal_assert(len(node.items) == 1, "Only one with element is supported so far!")
- context = node.items[0].context_expr
- option = node.items[0].optional_vars
- _internal_assert(isinstance(context, ast.Call), "The object must be a Python func call!")
- _internal_assert(isinstance(option, ast.Name), "The object after 'as' must be an id!")
- self.annotation[option.id] = context.func.id
- return visit_list_to_block(self.visit, node.body)
-
- def visit_If(self, node):
- cond = self.analyzer.simplify(self.visit(node.test))
-
- # Return no IfThenElse if proven
- if isinstance(cond, _expr.IntImm):
- if cond.value:
- return visit_list_to_block(self.visit, node.body)
- if node.orelse:
- return visit_list_to_block(self.visit, node.orelse)
- return utils.make_nop()
-
- if_body = visit_list_to_block(self.visit, node.body)
-
- if node.orelse:
- else_body = visit_list_to_block(self.visit, node.orelse)
- else:
- else_body = None
- return tvm.tir.IfThenElse(cond, if_body, else_body)
-
- def visit_IfExp(self, node):
- cond = self.visit(node.test)
- if_body = self.visit(node.body)
- else_body = self.visit(node.orelse)
- return tvm.tir.Select(cond, if_body, else_body)
-
- def visit_Compare(self, node):
- _internal_assert(len(node.ops) == len(node.comparators), "#compare ops != #comparators")
- ops = [self.visit(node.left)]
- ops += [self.visit(i) for i in node.comparators]
- res = []
- for i in range(len(node.ops)):
- lhs = ops[i]
- rhs = ops[i + 1]
- res.append(HybridParser._binop_maker[type(node.ops[i])](lhs, rhs))
- return _all(*res)
-
- def visit_BoolOp(self, node):
- n = len(node.values)
- if n == 1:
- _internal_assert(isinstance(node.op, ast.Not), "Unary is supposed to be not!")
- return operator.not_(self.visit(node.values[0]))
- _internal_assert(isinstance(node.op, (ast.And, ast.Or)), "Binary is supposed to be and/or!")
- values = [self.visit(i) for i in node.values]
- return HybridParser._binop_maker[type(node.op)](*values)
-
- def visit_UnaryOp(self, node):
- operand = self.visit(node.operand)
- return HybridParser._unaryop_maker[type(node.op)](operand)
-
- def visit_BinOp(self, node):
- lhs = self.visit(node.left)
- rhs = self.visit(node.right)
- return HybridParser._binop_maker[type(node.op)](lhs, rhs)
-
- def visit_Call(self, node):
- # Yet, no function pointer supported
- _internal_assert(
- isinstance(node.func, ast.Name), "Only id-function function call is supported so far!"
- )
-
- func_id = node.func.id
- args = [self.visit(i) for i in node.args]
- # Intrinsics'
- if hasattr(calls, func_id):
- return getattr(calls, func_id)(func_id, args)
- # Contexts'
- _internal_assert(
- func_id in self.symbols.keys(),
- f"The function called ({func_id}) is not in the context either!",
- )
- ty, entry = self.symbols[func_id]
- _internal_assert(ty is Symbol.Callable, "Are you sure what you call is a function?!")
- outs = entry(*args)
- op = outs.op if isinstance(outs, Tensor) else outs[0].op
- return op
-
- def visit_For(self, node):
- iter_var, low, ext, kind = self.visit(node.iter)
- _internal_assert(
- isinstance(node.target, ast.Name), "The loop iterator should be a variable!"
- )
-
- _name = node.target.id
-
- if isinstance(kind, tuple):
- low = self.analyzer.simplify(low)
- ext = self.analyzer.simplify(ext)
- _internal_assert(
- isinstance(low, _expr.ConstExpr) and isinstance(ext, _expr.ConstExpr),
- "Const range should start from a const " + "and iterate const times",
- )
-
- low, ext = low.value, ext.value
- if ext > 114514:
- logging.log(
- logging.CRITICAL, "[Warning] Are you sure to unroll a large loop in Python?"
- )
-
- bodies = []
- for i in range(low, low + ext):
- self.add_symbol(_name, Symbol.ConstLoopVar, i)
- body = visit_list_to_block(self.visit, node.body)
- body = self.wrap_up_realize(node, body)
- bodies.append(body)
- self.symbols.pop(_name)
- return concat_list_to_block(bodies)
-
- if iter_var is None:
- _internal_assert(kind is not None, "The loop iterating function parse error!")
- if isinstance(ext, _expr.PrimExpr):
- dtype = ext.dtype
- elif isinstance(ext, int):
- dtype = "int32"
- else:
- raise NotImplementedError(f"Unsupported type of ext: {type(ext)}")
- offset = iter_var = tvm.te.var(_name, dtype=dtype)
- if not tvm.tir.analysis.expr_deep_equal(low, tvm.runtime.const(0, "int32")):
- offset = iter_var + low
- self.add_symbol(_name, Symbol.LoopVar, offset)
- _body = visit_list_to_block(self.visit, node.body)
- else:
- _internal_assert(kind is None, "The loop bind function parse error!")
- self.add_symbol(_name, Symbol.ThreadBind, iter_var)
- self.device += 1
- _body = visit_list_to_block(self.visit, node.body)
- self.device -= 1
-
- _body = self.wrap_up_realize(node, _body)
-
- if kind is None:
- res = _body
- else:
- _internal_assert(
- not isinstance(kind, tuple), "Micro expansion should be handled before!"
- )
- res = tvm.tir.For(iter_var, tvm.runtime.const(0, "int32"), ext, kind, _body)
-
- self.symbols.pop(_name)
- return res
-
- def visit_Return(self, node):
- _internal_assert(
- all(ty != Symbol.LoopVar for ty, _ in self.symbols.values()),
- "Return should not be in a loop body!",
- )
- ids = []
- if isinstance(node.value, ast.Name):
- ids = [node.value.id]
- else:
- _internal_assert(
- isinstance(node.value, ast.Tuple),
- "You should return either a single tensor or a tuple",
- )
- _internal_assert(
- all(isinstance(i, ast.Name) for i in node.value.elts), "What do you return?"
- )
- ids = [i.id for i in node.value.elts]
- _internal_assert(len(set(ids)) == len(ids), "Duplicated tensors in the return tuples")
- if len(ids) < len(self.outputs):
- logging.log(logging.CRITICAL, "[Warning] Not all the output buffers returned!")
- self.outputs = [self.symbols[i][1] for i in ids]
- self.returned = True
- return utils.make_nop()
-
- def visit_Tuple(self, node):
- return tuple(self.visit(i) for i in node.elts)
-
- def visit_Str(self, node):
- return node.s
-
- def visit_Assert(self, node):
- test = self.visit(node.test)
- mesg = tvm.runtime.convert(self.visit(node.msg))
- return tvm.tir.AssertStmt(test, mesg, utils.make_nop())
-
-
-def parse_python(src, args, symbols, closure_vars):
- """The helper function of calling the AST visitor
-
- Parameters
- ----------
- src : ast.node or str
- If an ast.node, then directly lower it.
- If a str, then parse it to ast and lower it.
-
- args : list of Tensors or Vars
- The argument lists to the function.
- It is NOT encouraged to write a function without arguments.
- It is NOT encouraged to write a function with side effect.
-
- symbols : list of str
- The symbol list of the global context of the function.
-
- closure_vars: dict
- A dict of external name reference captured by this function.
-
- Returns
- -------
- root : Stmt
- The result Halide IR and the parser class instance.
- """
- root = ast.parse(src) if isinstance(src, str) else src
- _internal_assert(root, ast.AST)
- var_usage = determine_variable_usage(root, args, symbols, closure_vars)
- parser = HybridParser(args, var_usage, symbols, closure_vars)
- parser.parsed_body = parser.visit(root)
- _internal_assert(parser.returned, "No valid return found in the function body!")
- return parser
-
-
-def source_to_op(src, args, symbols, closure_vars):
- """Another level of wrapper
-
- Parameters
- ----------
- src : ast.node or str
- If an ast.node, then directly lower it.
- If a str, then parse it to ast and lower it.
-
- args : list of Tensors or Vars
- The argument lists to the function.
- It is NOT encouraged to write a function without arguments.
- It is NOT encouraged to write a function with side effect.
-
- symbols : list of str
- The symbol list of the global context of the function.
-
- closure_vars: dict
- A dict of external name reference captured by this function.
-
- Returns
- -------
- res : list of output tensors
- The result of output tensors of the formed OpNode.
- """
- parser = parse_python(src, args, symbols, closure_vars)
-
- input_tensors = []
-
- def get_input_tensors(arg):
- if isinstance(arg, Tensor):
- input_tensors.append(arg)
- elif isinstance(arg, Array):
- for i in arg:
- get_input_tensors(i)
-
- for i in args:
- get_input_tensors(i)
- op = tvm.te._ffi_api.HybridOp(
- parser.func_name, "HybridOp", None, input_tensors, parser.outputs, parser.parsed_body
- )
- res = [op.output(i) for i in range(len(parser.outputs))]
- return res[0] if len(res) == 1 else res
diff --git a/python/tvm/te/hybrid/preprocessor.py b/python/tvm/te/hybrid/preprocessor.py
deleted file mode 100644
index 6af584060e..0000000000
--- a/python/tvm/te/hybrid/preprocessor.py
+++ /dev/null
@@ -1,120 +0,0 @@
-# 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.
-"""Determines the declaration, r/w status, and last use of each variable"""
-
-import ast
-import sys
-from .runtime import HYBRID_GLOBALS
-from .utils import _internal_assert
-
-
-class PyVariableUsage(ast.NodeVisitor):
- """The vistor class to determine the declaration, r/w status, and last use of each variable"""
-
- # pylint: disable=invalid-name
- # pylint: disable=missing-docstring
- def __init__(self, args, symbols, closure_vars):
- self.status = {}
- self.scope_level = []
- self._args = {}
- self.args = args
- self.aug_assign_ = False
- self.symbols = symbols
- self.closure_vars = closure_vars
-
- def visit_FunctionDef(self, node):
- self.scope_level.append(node)
- _internal_assert(
- len(node.args.args) == len(self.args),
- "#arguments passed should be the same as #arguments defined",
- )
- for idx, arg in enumerate(node.args.args):
- _attr = "id" if sys.version_info[0] < 3 else "arg" # To make py2 and 3 compatible
- self._args[getattr(arg, _attr)] = self.args[idx]
- for i in node.body:
- self.visit(i)
-
- def visit_For(self, node):
- _internal_assert(isinstance(node.target, ast.Name), "For's iterator should be an id")
- self.visit(node.iter)
- self.scope_level.append(node)
- for i in node.body:
- self.visit(i)
- self.scope_level.pop()
-
- def visit_Call(self, node):
- # No function pointer supported so far
- _internal_assert(isinstance(node.func, ast.Name), "Function call should be an id")
- func_id = node.func.id
- _internal_assert(
- func_id
- in list(HYBRID_GLOBALS.keys())
- + ["range", "max", "min", "len"]
- + list(self.symbols.keys()),
- "Function call id " + func_id + " not in intrinsics' list",
- )
- for elem in node.args:
- self.visit(elem)
-
- def visit_AugAssign(self, node):
- self.aug_assign_ = True
- self.generic_visit(node)
- self.aug_assign_ = False
-
- def visit_Name(self, node):
- # If it is True or False, we do not worry about it!
- if sys.version_info[0] == 2 and node.id in ["True", "False"]:
- return
- # If it is from the argument list or loop variable, we do not worry about it!
- if node.id in self._args.keys():
- return
- fors = [loop.target.id for loop in self.scope_level if isinstance(loop, ast.For)]
- if node.id in fors:
- return
- # The loop variable cannot be overwritten when iteration
- _internal_assert(
- not isinstance(node.ctx, ast.Store) or node.id not in fors,
- "Iter var cannot be overwritten",
- )
-
- if node.id not in self.status.keys():
- # It is a captured value in closure
- if node.id in self.closure_vars:
- try:
- ast.literal_eval(str(self.closure_vars[node.id]))
- except ValueError:
- raise ValueError("Only support capturing constant values in closure")
- return
-
- _internal_assert(isinstance(node.ctx, ast.Store), f"Undeclared variable {node.id}")
- if self.aug_assign_:
- raise ValueError('"First store" cannot be an AugAssign')
- self.status[node.id] = (node, self.scope_level[-1], set())
- else:
- decl, loop, usage = self.status[node.id]
- usage.add(type(node.ctx))
- _internal_assert(
- loop in self.scope_level, f"{node.id} is used out of the scope it is defined!"
- )
- self.status[node.id] = (decl, loop, usage)
-
-
-def determine_variable_usage(root, args, symbols, closure_vars):
- """The helper function for calling the dedicated visitor."""
- visitor = PyVariableUsage(args, symbols, closure_vars)
- visitor.visit(root)
- return visitor.status
diff --git a/python/tvm/te/hybrid/runtime.py b/python/tvm/te/hybrid/runtime.py
deleted file mode 100644
index 615bd7e43a..0000000000
--- a/python/tvm/te/hybrid/runtime.py
+++ /dev/null
@@ -1,175 +0,0 @@
-# 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.
-"""Intrinsics of TVM-Python Hybrid Script for Python emulation runtime"""
-
-import numpy
-from tvm.target import Target
-
-
-class bind(object): # pylint: disable=invalid-name
- """GPU bind software emulataion runtime."""
-
- def __init__(self, _, ext):
- self.ext = ext
-
- def __iter__(self):
- i = 0
- while i < self.ext:
- yield i
- i += 1
-
-
-def allocate(shape, dtype="float32", scope="global"): # pylint: disable=unused-argument
- """Allocate a buffer with given shape
-
- Parameters
- ----------
- shape: Tuple
- The shape of the tensor to be allocated
- dtype: string
- The data type of the tensor
- scope: string
- The storage scope of the tensor
-
- Returns
- -------
- tensor: numpy.array
- The tensor allocated
- """
- return numpy.zeros(shape).astype(dtype)
-
-
-def rsqrt(x):
- """
- Computes reciprocal of square root of x element-wise
-
- Parameters
- ----------
- x: Tensor
-
- Returns
- -------
- res: Tensor
- The result of reciprocal of square root of x
- """
- return numpy.ones_like(x) / numpy.sqrt(x)
-
-
-def popcount(x):
- """
- Count ones in the binary representation of number x
-
- Parameters
- ----------
- x: Integer
- The number to be counted
-
- Returns
- -------
- cnt: Integer
- The number of ones in the binary representation of number x
- """
- cnt = 0
- while x:
- x -= x & -x
- cnt += 1
- return cnt
-
-
-def sigmoid(x):
- """
- Sigmoid function of x, aka 1/(1+exp(-x)).
-
- Parameters
- ----------
- x: a real number
-
- Returns
- -------
- res: a real number
- The result of sigmoid function
- """
- return 1 / (1 + numpy.exp(-x))
-
-
-def max_num_threads(allow_none=True):
- """Get max number of threads for GPU targets."""
- return Target.current(allow_none).max_num_threads
-
-
-def inf(dtype):
- return numpy.iinfo(dtype).max
-
-
-def ninf(dtype):
- return numpy.iinfo(dtype).min
-
-
-HYBRID_GLOBALS = {
- "unroll": range,
- "vectorize": range,
- "parallel": range,
- "const_range": range,
- "bind": bind,
- "allocate": allocate,
- "output_tensor": allocate,
- "sqrt": numpy.sqrt,
- "rsqrt": rsqrt,
- "log": numpy.log,
- "tanh": numpy.tanh,
- "power": numpy.power,
- "exp": numpy.exp,
- "sigmoid": sigmoid,
- "popcount": popcount,
- "round": round,
- "likely": lambda cond: cond,
- "uint8": numpy.uint8,
- "uint16": numpy.uint16,
- "uint32": numpy.uint32,
- "uint64": numpy.uint64,
- "int8": numpy.int8,
- "int16": numpy.int16,
- "int32": numpy.int32,
- "int64": numpy.int64,
- "float16": numpy.float16,
- "float32": numpy.float32,
- "float64": numpy.float64,
- "ceil_div": lambda a, b: (a + b - 1) // b,
- "max_num_threads": max_num_threads,
- "inf": inf,
- "ninf": inf,
-}
-
-
-def _enter_hybrid_runtime(func):
- """Put hybrid runtime variables into the global scope"""
- _globals = func.__globals__
- intersect = []
- for elem in list(HYBRID_GLOBALS.keys()):
- if elem in _globals.keys():
- intersect.append((elem, _globals[elem]))
- _globals[elem] = HYBRID_GLOBALS[elem]
- return intersect
-
-
-def _restore_runtime(func, intersect):
- """Rollback the modification caused by hybrid runtime"""
- _globals = func.__globals__
- for elem in list(HYBRID_GLOBALS.keys()):
- _globals.pop(elem)
- for k, v in intersect:
- _globals[k] = v
diff --git a/python/tvm/te/hybrid/utils.py b/python/tvm/te/hybrid/utils.py
deleted file mode 100644
index a515938fa5..0000000000
--- a/python/tvm/te/hybrid/utils.py
+++ /dev/null
@@ -1,103 +0,0 @@
-# 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.
-# pylint: disable=inconsistent-return-statements
-"""Internal utilities for parsing Python subset to TIR"""
-
-import ast
-import inspect
-import logging
-import sys
-import numpy
-
-import tvm.runtime
-from tvm._ffi.base import numeric_types
-from tvm.ir.container import Array
-
-from tvm.tir import expr as _expr
-from tvm.tir import stmt as _stmt
-from tvm.te.tensor import Tensor
-
-
-# pylint: disable=invalid-name
-np_arg_types = (numpy.ndarray, *numeric_types)
-tvm_arg_types = (Tensor, Array, _expr.Var, _expr.ConstExpr, *numeric_types, list, tuple, str)
-halide_imm_types = (_expr.IntImm, _expr.FloatImm, *numeric_types)
-
-
-def _internal_assert(cond, err):
- """Simplify the code segment like if not XXX then raise an error"""
- if not cond:
- raise ValueError(err)
-
-
-# Useful constants. In avoid of runtime dependences, we use function calls to return them.
-def make_nop():
- """Returns a 'no operation' node in HalideIR."""
- return _stmt.Evaluate(tvm.runtime.const(0, dtype="int32"))
-
-
-def is_docstring(node):
- """Checks if a Python AST node is a docstring"""
- return isinstance(node, ast.Expr) and isinstance(node.value, ast.Str)
-
-
-def _pruned_source(func):
- """Prune source code's extra leading spaces"""
- try:
- lines = inspect.getsource(func).split("\n")
- leading_space = len(lines[0]) - len(lines[0].lstrip(" "))
- lines = [line[leading_space:] for line in lines]
- return "\n".join(lines)
- except IOError as err:
- if sys.version_info[0] == 2 and str(err) == "could not get source code":
- logging.log(
- logging.CRITICAL,
- "This module is not fully operated under Python2... " "Please move to Python3!",
- )
- raise err
-
-
-def replace_io(body, rmap):
- """Replacing tensors usage according to the dict given"""
- # pylint: disable=import-outside-toplevel
- from tvm.tir import stmt_functor
-
- def replace(op):
- if isinstance(op, _stmt.ProducerStore) and op.producer.op in rmap.keys():
- buf = rmap[op.producer.op]
- return _stmt.ProducerStore(buf, op.value, op.indices)
- if isinstance(op, _expr.ProducerLoad) and op.producer.op in rmap.keys():
- buf = rmap[op.producer.op]
- return _expr.ProducerLoad(buf, op.indices)
- return None
-
- return stmt_functor.ir_transform(body, None, replace, ["tir.ProducerStore", "tir.ProducerLoad"])
-
-
-def _is_tvm_arg_types(args):
- """Determine a list of element is either a list of tvm arguments of a list of numpy arguments.
- If neither is true, raise a value error."""
- if all(isinstance(elem, tvm_arg_types) for elem in args):
- return True
- elif all(isinstance(elem, np_arg_types) for elem in args):
- return False
- else:
- raise ValueError(
- f"Expected arguments to be entirely TVM types, "
- f"or entirely numpy types, "
- f"but received {[type(elem) for elem in args]}"
- )
diff --git a/python/tvm/te/schedule.py b/python/tvm/te/schedule.py
deleted file mode 100644
index 87a4eda728..0000000000
--- a/python/tvm/te/schedule.py
+++ /dev/null
@@ -1,665 +0,0 @@
-# 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.
-# pylint: disable=unused-import
-"""The computation schedule api of TVM."""
-import collections
-import inspect
-from typing import Callable, List
-
-import tvm._ffi
-from tvm._ffi.base import string_types
-from tvm.ir import container as _container
-from tvm.runtime import Object, convert
-from tvm.tir import Buffer, IndexMap, IterVar, Var
-
-from . import _ffi_api
-from . import tensor as _tensor
-
-
-@tvm._ffi.register_object
-class Split(Object):
- """Split operation on axis."""
-
-
-@tvm._ffi.register_object
-class Fuse(Object):
- """Fuse operation on axis."""
-
-
-@tvm._ffi.register_object
-class Singleton(Object):
- """Singleton axis."""
-
-
-def create_schedule(ops):
- """Create a schedule for list of ops
-
- Parameters
- ----------
- ops : list of Operations
- The source expression.
-
- Returns
- -------
- sch : schedule.Schedule
- The created schedule.
- """
- if not isinstance(ops, (list, _container.Array)):
- ops = [ops]
- return _ffi_api.CreateSchedule(ops)
-
-
-@tvm._ffi.register_object
-class Schedule(Object):
- """Schedule for all the stages."""
-
- def __getitem__(self, k):
- if isinstance(k, _tensor.Tensor):
- k = k.op
- if not isinstance(k, _tensor.Operation):
- raise ValueError("Expect schedule key to be Tensor or Operation")
- if k not in self.stage_map:
- raise ValueError(f"Cannot find the operation {k} in schedule")
- return self.stage_map[k]
-
- def normalize(self):
- """Build a normalized schedule from the current schedule.
-
- Insert necessary rebase to make certain iter var to start from 0.
- This is needed before bound inference and followup step.
-
- Returns
- -------
- sch : Schedule
- The normalized schedule.
- """
- return _ffi_api.ScheduleNormalize(self)
-
- def create_group(self, outputs, inputs, include_inputs=False):
- """Create stage group by giving output and input boundary.
-
- The operators between outputs and inputs are placed as member of group.
- outputs are include in the group, while inputs are not included.
-
- Parameters
- ----------
- outputs : list of Tensors
- The outputs of the group.
-
- inputs : list of Tensors
- The inputs of the group.
-
- include_inputs : boolean, optional
- Whether include input operations in the group if they are used by outputs.
-
- Returns
- -------
- group : Stage
- A virtual stage represents the group, user can use compute_at to move
- the attachment point of the group.
- """
- if isinstance(outputs, _tensor.Tensor):
- outputs = [outputs]
- if isinstance(inputs, _tensor.Tensor):
- inputs = [inputs]
- return _ffi_api.ScheduleCreateGroup(self, outputs, inputs, include_inputs)
-
- def cache_read(self, tensor, scope, readers):
- """Create a cache read of original tensor for readers.
-
- This will mutate the body of the readers.
- A new cache stage will be created for the tensor.
- Call this before doing any split/fuse schedule.
-
- Parameters
- ----------
- tensor : Tensor
- The tensor to be cached.
- scope : str
- The scope of cached
- readers : list of Tensor or Operation
- The readers to read the cache.
-
- Returns
- -------
- cache : Tensor
- The created cache tensor.
- """
- if isinstance(readers, (_tensor.Tensor, _tensor.Operation)):
- readers = [readers]
- readers = [t.op if isinstance(t, _tensor.Tensor) else t for t in readers]
- return _ffi_api.ScheduleCacheRead(self, tensor, scope, readers)
-
- def cache_write(self, tensor, scope):
- """Create a cache write of original tensor, before storing into tensor.
-
- This will mutate the body of the tensor.
- A new cache stage will created before feed into the tensor.
-
- This function can be used to support data layout transformation.
- If there is a split/fuse/reorder on the data parallel axis of tensor
- before cache_write is called. The intermediate cache stores
- the data in the layout as the iteration order of leave axis.
- The data will be transformed back to the original layout in the original tensor.
- User can further call compute_inline to inline the original layout and keep
- the data stored in the transformed layout.
-
- Parameters
- ----------
- tensor : Tensor, list or tuple
- The tensors to be feed to. All the tensors must be produced by one computeOp
- scope : str
- The scope of cached
-
- Returns
- -------
- cache : Tensor
- The created cache tensor.
- """
- return _ffi_api.ScheduleCacheWrite(self, tensor, scope)
-
- def rfactor(self, tensor, axis, factor_axis=0):
- """Factor a reduction axis in tensor's schedule to be an explicit axis.
-
- This will create a new stage that generated the new tensor with axis
- as the first dimension. The tensor's body will be rewritten as a reduction
- over the factored tensor.
-
- Parameters
- ----------
- tensor : Tensor
- The tensor to be factored.
- axis : IterVar
- The reduction axis in the schedule to be factored.
- factor_axis : int
- The position where the new axis is placed.
-
- Returns
- -------
- tfactor : Tensor or Array of Tensor
- The created factored tensor.
- """
- factored = _ffi_api.ScheduleRFactor(self, tensor, axis, factor_axis)
- return factored[0] if len(factored) == 1 else factored
-
-
-@tvm._ffi.register_object
-class Stage(Object):
- """A Stage represents schedule for one operation."""
-
- def split(self, parent, factor=None, nparts=None, disable_predication=False):
- """Split the stage either by factor providing outer scope, or both
-
- Parameters
- ----------
- parent : IterVar
- The parent iter var.
-
- factor : Expr, optional
- The splitting factor
-
- nparts : Expr, optional
- The number of outer parts.
-
- disable_predication : bool, optional
- If enabled, don't create a predicate for guarding the loop. This can
- be useful when splitting with scalable factors that the schedule writer
- knows are divisible by the loop bound.
-
- Warning: enabling this feature may result in incorrect code generation
- if not used carefully.
-
- Returns
- -------
- outer : IterVar
- The outer variable of iteration.
-
- inner : IterVar
- The inner variable of iteration.
- """
- if nparts is not None:
- if factor is not None:
- raise ValueError("Do not need to provide both outer and nparts")
- outer, inner = _ffi_api.StageSplitByNParts(self, parent, nparts, disable_predication)
- else:
- if factor is None:
- raise ValueError("Either nparts or factor need to be provided")
- outer, inner = _ffi_api.StageSplitByFactor(self, parent, factor, disable_predication)
- return outer, inner
-
- def fuse(self, *args):
- """Fuse multiple consecutive iteration variables into a single iteration variable.
-
- fused = fuse(...fuse(fuse(args[0], args[1]), args[2]),..., args[-1])
- The order is from outer to inner.
-
- Parameters
- ----------
- args : list of IterVars
- Itervars that proceeds each other
-
- Returns
- -------
- fused : IterVar
- The fused variable of iteration.
- """
- fused = _ffi_api.StageFuse(self, args)
- return fused
-
- def set_scope(self, scope):
- """Set the thread scope of this stage
-
- Parameters
- ----------
- scope : str
- The thread scope of this stage
- """
- return _ffi_api.StageSetScope(self, scope)
-
- def bind(self, ivar, thread_ivar):
- """Bind ivar to thread index thread_ivar
-
- Parameters
- ----------
- ivar : IterVar
- The iteration to be binded to thread.
-
- thread_ivar : IterVar
- The thread to be binded.
- """
- _ffi_api.StageBind(self, ivar, thread_ivar)
-
- def env_threads(self, threads):
- """Mark threads to be launched at the outer scope of composed op.
-
- Parameters
- ----------
- threads : list of threads
- The threads to be launched.
- """
- if isinstance(threads, IterVar):
- threads = [threads]
- _ffi_api.StageEnvThreads(self, threads)
-
- def set_store_predicate(self, predicate):
- """Set predicate under which store to the array can be performed.
-
- Use this when there are duplicated threads doing the same store and we only
- need one of them to do the store.
-
- Parameters
- ----------
- predicate : Expr
- The guard condition fo store.
- """
- _ffi_api.StageSetStorePredicate(self, predicate)
-
- def compute_at(self, parent, scope):
- """Attach the stage at parent's scope
-
- Parameters
- ----------
- parent : Stage
- The parent stage
-
- scope : IterVar
- The loop scope t be attached to.
- """
- _ffi_api.StageComputeAt(self, parent, scope)
-
- def compute_inline(self):
- """Mark stage as inline
-
- Parameters
- ----------
- parent : Stage
- The parent stage
- """
- _ffi_api.StageComputeInline(self)
-
- def compute_root(self):
- """Attach the stage at parent, and mark it as root
-
- Parameters
- ----------
- parent : Stage
- The parent stage
- """
- _ffi_api.StageComputeRoot(self)
-
- def reorder(self, *args):
- """reorder the arguments in the specified order.
-
- Parameters
- ----------
- args : list of IterVar
- The order to be ordered
- """
- _ffi_api.StageReorder(self, args)
-
- def tile(self, x_parent, y_parent, x_factor, y_factor):
- """Perform tiling on two dimensions
-
- The final loop order from outmost to inner most are
- [x_outer, y_outer, x_inner, y_inner]
-
- Parameters
- ----------
- x_parent : IterVar
- The original x dimension
- y_parent : IterVar
- The original y dimension
- x_factor : Expr
- The stride factor on x axis
- y_factor : Expr
- The stride factor on y axis
-
- Returns
- -------
- x_outer : IterVar
- Outer axis of x dimension
- y_outer : IterVar
- Outer axis of y dimension
- x_inner : IterVar
- Inner axis of x dimension
- p_y_inner : IterVar
- Inner axis of y dimension
- """
- x_outer, y_outer, x_inner, y_inner = _ffi_api.StageTile(
- self, x_parent, y_parent, x_factor, y_factor
- )
- return x_outer, y_outer, x_inner, y_inner
-
- def vectorize(self, var):
- """Vectorize the iteration.
-
- Parameters
- ----------
- var : IterVar
- The iteration to be vectorize
- """
- _ffi_api.StageVectorize(self, var)
-
- def tensorize(self, var, tensor_intrin):
- """Tensorize the computation enclosed by var with tensor_intrin
-
- Parameters
- ----------
- var : IterVar
- The iteration boundary of tensorization.
-
- tensor_intrin : TensorIntrin
- The tensor intrinsic used for computation.
- """
- _ffi_api.StageTensorize(self, var, tensor_intrin)
-
- def unroll(self, var):
- """Unroll the iteration.
-
- Parameters
- ----------
- var : IterVar
- The iteration to be unrolled.
- """
- _ffi_api.StageUnroll(self, var)
-
- def parallel(self, var):
- """Parallelize the iteration.
-
- Parameters
- ----------
- var : IterVar
- The iteration to be parallelized.
- """
- _ffi_api.StageParallel(self, var)
-
- def pragma(self, var, pragma_type, pragma_value=None):
- """Annotate the iteration with pragma
-
- This will translate to a pragma_scope surrounding
- the corresponding loop generated.
- Useful to support experimental features and extensions.
-
- Parameters
- ----------
- var : IterVar
- The iteration to be anotated
-
- pragma_type : str
- The pragma string to be annotated
-
- pragma_value : Expr, optional
- The pragma value to pass along the pragma
-
- Note
- ----
- Most pragmas are advanced/experimental features
- and may subject to change. List of supported pragmas:
-
- - **debug_skip_region**
-
- Force skip the region marked by the axis and turn it into no-op.
- This is useful for debug purposes.
-
- - **parallel_launch_point**
-
- Specify to launch parallel threads outside the
- specified iteration loop. By default the threads
- launch at the point of parallel construct.
- This pragma moves the launching point to even outer scope.
- The threads are launched once and reused across multiple
- parallel constructs as BSP style program.
-
- - **parallel_barrier_when_finish**
-
- Insert a synchronization barrier between working threads
- after the specified loop iteration finishes.
-
- - **parallel_stride_pattern**
-
- Hint parallel loop to execute in strided pattern.
- :code:`for (int i = task_id; i < end; i += num_task)`
-
- """
- if isinstance(pragma_value, string_types):
- pragma_value = convert(pragma_value)
- _ffi_api.StagePragma(self, var, pragma_type, pragma_value)
-
- def prefetch(self, tensor, var, offset):
- """Prefetch the specified variable
-
- Parameters
- ----------
- tensor : Tensor
- The tensor to be prefetched
- var : IterVar
- The loop point at which the prefetching is applied
- offset : Expr
- The number of iterations to be prefetched before actual execution
- """
- _ffi_api.StagePrefetch(self, tensor, var, offset)
-
- def storage_align(self, axis, factor, offset):
- """Set alignment requirement for specific axis
-
- This ensures that stride[axis] == k * factor + offset for some k.
- This is useful to set memory layout to for more friendly memory
- access pattern. For example, we can set alignment to be
- factor=2, offset=1 to avoid bank conflict for thread access on
- higher dimension in GPU shared memory.
-
- Parameters
- ----------
- axis : IterVar
- The axis dimension to be aligned.
- factor : int
- The factor in alignment specification.
- offset : int
- The offset in the alignment specification.
- """
- _ffi_api.StageStorageAlign(self, axis, factor, offset)
-
- def double_buffer(self):
- """Compute the current stage via double buffering.
-
- This can only be applied to intermediate stage.
- This will double the storage cost of the current stage.
- Can be useful to hide load latency.
- """
- _ffi_api.StageDoubleBuffer(self)
-
- def rolling_buffer(self):
- """Compute the current stage via rolling buffering.
-
- This can only be applied to intermediate stage.
- This will change the storage cost of the current stage.
- """
- _ffi_api.StageRollingBuffer(self)
-
- def transform_layout(self, mapping_function: Callable[..., List[tvm.tir.PrimExpr]]):
- """Defines the layout transformation for the current stage's tensor.
-
- The map from initial_indices to final_indices must be an
- invertible affine transformation. This method may be called
- more than once for a given tensor, in which case each
- transformation is applied sequentially.
-
- If the stage is a ComputeOp, then the iteration order of the
- compute stage is rewritten to be a row-major traversal of the
- tensor, and the new loop iteration variables are returned.
- For all other stages, the loop iteration order is unmodified,
- and the return value is None.
-
- Parameters
- ----------
- mapping_function : Callable[..., List[tvm.tir.PrimExpr]]
-
- A callable that accepts N arguments of type tvm.tir.Var,
- and outputs a list of PrimExpr. The input arguments
- represent the location of a value in the current stage's
- tensor, using the pre-transformation layout. The return
- value of the function gives the location of that value in
- the current stage's tensor, using the post-transformation
- layout.
-
- Returns
- -------
- new_iter_vars : Optional[List[tvm.tir.IterVar]]
-
- If the stage is a ComputeOp, then the return will be the
- updated loop iteration variables over the data array, in
- the same order as the output values from the
- `mapping_function`.
-
- Otherwise, the return value is None.
-
- Examples
- --------
- .. code-block:: python
-
- # ``A`` is a tensor whose compute definition is in NHWC
- # format, and should be transformed into NCHWc format.
-
- s[A].transform_layout(
- lambda n,h,w,c: [n, c//4, h, w, c%4]
- )
-
-
- .. code-block:: python
-
- # ``A`` is a tensor whose compute definition is in an
- # arbitrary format, and should be transformed such that
- # the last index is split, with the slower-changing index
- # of the split placed at the slowest changing dimension.
-
- s[A].transform_layout(
- lambda *indices, i: [i//4, *indices, i%4]
- )
-
- .. code-block:: python
-
- # ``B`` is a tensor defined by te.compute to be a copy of
- # ``A`, and should be transformed such that ``B``'s layout
- # is a transpose of ``A``'s layout. The loop iteration
- # that computes ``B`` will correspond to ``B``'s memory
- # layout.
-
- A = te.placeholder([n,m])
- B = te.compute(A.shape, lambda i,j: A[i,j])
- s = te.create_schedule(B.op)
-
- s[B].transform_layout(lambda i,j: [j,i])
-
- """
-
- ndim = len(self.op.output(0).shape)
- index_map, axis_separators = IndexMap.from_func_with_separators(
- mapping_function, ndim=ndim, index_dtype="int32"
- )
-
- new_iter_vars = _ffi_api.StageTransformLayout(
- self, index_map.initial_indices, index_map.final_indices
- )
- _ffi_api.StageSetAxisSeparators(self, axis_separators)
-
- return new_iter_vars or None
-
-
-@tvm._ffi.register_object
-class SpecializedCondition(Object):
-
- """Specialized condition to enable op specialization."""
-
- def __init__(self, conditions):
- """Create a specialized condition.
-
- .. note::
- Conditions are represented in conjunctive joint form (CNF).
- Each condition should be a simple expression, e.g., n > 16,
- m % 8 == 0, etc., where n, m are tvm.Var that represents a
- dimension in the tensor shape.
-
- Parameters
- ----------
- conditions : List of tvm.Expr
- List of conditions in conjunctive joint form (CNF).
- """
- if not isinstance(conditions, (list, _container.Array)):
- conditions = [conditions]
- self.__init_handle_by_constructor__(_ffi_api.CreateSpecializedCondition, conditions)
-
- @staticmethod
- def current():
- """Returns the current specialized condition"""
- return _ffi_api.GetCurrentSpecialization()
-
- def __enter__(self):
- _ffi_api.EnterSpecializationScope(self)
- return self
-
- def __exit__(self, ptype, value, trace):
- _ffi_api.ExitSpecializationScope(self)
-
-
-# Sentinel value used to indicate which groups of pre-flattening axes
-# should be used to post-flattening axes. Moved from
-# te.AXIS_SEPARATOR to tir.IndexMap.AXIS_SEPARATOR for general use,
-# maintained here for backwards compatibility.
-AXIS_SEPARATOR = IndexMap.AXIS_SEPARATOR
-
-
-tvm._ffi._init_api("schedule", __name__)
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index 930667242e..53ab9d0b5b 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -190,13 +190,3 @@ class ScanOp(Operation):
@tvm._ffi.register_object
class ExternOp(Operation):
"""External operation."""
-
-
-@tvm._ffi.register_object
-class HybridOp(Operation):
- """Hybrid operation."""
-
- @property
- def axis(self):
- """Represent the IterVar axis, also defined when it is a HybridOp"""
- return self.__getattr__("axis")
diff --git a/python/tvm/te/tensor_intrin.py b/python/tvm/te/tensor_intrin.py
deleted file mode 100644
index ff633af02d..0000000000
--- a/python/tvm/te/tensor_intrin.py
+++ /dev/null
@@ -1,146 +0,0 @@
-# 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.
-"""Tensor intrinsics"""
-import tvm._ffi
-import tvm.tir
-
-from tvm.runtime import Object, convert
-from tvm.ir import Range
-from .tensor import PlaceholderOp
-
-from . import tensor as _tensor
-from . import _ffi_api
-
-
-def _get_region(tslice):
- region = []
- for idx in tslice.indices:
- if isinstance(idx, slice):
- assert idx.step is None
- region.append(Range(idx.start, idx.stop))
- else:
- if isinstance(idx, tvm.tir.IterVar):
- begin = idx.var
- else:
- begin = idx
- region.append(Range.from_min_extent(begin, 1))
- return region
-
-
-@tvm._ffi.register_object
-class TensorIntrin(Object):
- """Tensor intrinsic functions for certain computation.
-
- See Also
- --------
- decl_tensor_intrin: Construct a TensorIntrin
- """
-
- def __call__(self, *args, **kwargs):
- tensors = [x.tensor for x in args if isinstance(x, _tensor.TensorSlice)]
- scalar_inputs = [x for x in args if not isinstance(x, _tensor.TensorSlice)]
- regions = [_get_region(x) for x in args if isinstance(x, _tensor.TensorSlice)]
- reduce_axis = []
- if "reduce_axis" in kwargs:
- reduce_axis = kwargs["reduce_axis"]
- if not isinstance(reduce_axis, (list, tuple)):
- reduce_axis = [reduce_axis]
- reduce_axis = convert(reduce_axis)
- if scalar_inputs:
- scalar_inputs = convert(scalar_inputs)
- return _ffi_api.TensorIntrinCall(self, tensors, regions, reduce_axis, scalar_inputs)
-
-
-def decl_tensor_intrin(
- op, fcompute, name="tensor_intrin", binds=None, scalar_params=None, default_buffer_params=None
-):
- """Declare a tensor intrinsic function.
-
- Parameters
- ----------
- op: Operation
- The symbolic description of the intrinsic operation
-
- fcompute: lambda function of inputs, outputs-> stmt
- Specifies the IR statement to do the computation.
- See the following note for function signature of fcompute
-
- .. note::
- **Parameters**
-
- - **ins** (list of :any:`tvm.tir.Buffer`) - Placeholder for each inputs
- - **outs** (list of :any:`tvm.tir.Buffer`) - Placeholder for each outputs
-
- **Returns**
-
- - **stmt** (:any:`tvm.tir.Stmt`, or tuple of three stmts)
- - If a single stmt is returned, it represents the body
- - If tuple of three stmts are returned they corresponds to body,
- reduce_init, reduce_update
-
- name: str, optional
- The name of the intrinsic.
-
- binds: dict of :any:`Tensor` to :any:`tvm.tir.Buffer`, optional
- Dictionary that maps the Tensor to Buffer which specified the data layout
- requirement of the function. By default, a new compact buffer is created
- for each tensor in the argument.
-
- scalar_params: a list of variables used by op, whose values will be passed
- as scalar_inputs when the tensor intrinsic is called.
-
- default_buffer_params: Optional[dict]
- Dictionary of buffer arguments to be passed when constructing a buffer.
-
- Returns
- -------
- intrin: TensorIntrin
- A TensorIntrin that can be used in tensorize schedule.
- """
- if not isinstance(op, _tensor.Operation):
- raise TypeError("expect Operation")
- inputs = op.input_tensors
- binds = binds if binds else {}
- tensors = list(inputs)
- for i in range(op.num_outputs):
- tensors.append(op.output(i))
-
- binds_list = []
- for t in inputs:
- if not isinstance(t.op, PlaceholderOp):
- raise ValueError("Do not yet support composition op")
-
- default_buffer_params = {} if default_buffer_params is None else default_buffer_params
- for t in tensors:
- buf = (
- binds[t]
- if t in binds
- else tvm.tir.decl_buffer(t.shape, t.dtype, t.op.name, **default_buffer_params)
- )
- binds_list.append(buf)
-
- if scalar_params:
- body = fcompute(binds_list[: len(inputs)], binds_list[len(inputs) :], scalar_params)
- else:
- body = fcompute(binds_list[: len(inputs)], binds_list[len(inputs) :])
- scalar_params = []
- if isinstance(body, (tvm.tir.PrimExpr, tvm.tir.Stmt)):
- body = [body]
- body = [tvm.tir.Evaluate(x) if isinstance(x, tvm.tir.PrimExpr) else x for x in body]
- if len(body) < 3:
- body += [None] * (3 - len(body))
- return _ffi_api.TensorIntrin(name, op, inputs, binds_list, scalar_params, *body)
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index 8df32c8105..b3123a20d3 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -327,8 +327,7 @@ def check_bool_expr_is_true(bool_expr, vranges, cond=None):
A = tvm.te.compute([r.extent.value for v, r in vranges.items()], _compute_body)
args = [tvm.nd.empty(A.shape, A.dtype)]
- sch = tvm.te.create_schedule(A.op)
- mod = tvm.build(sch, [A])
+ mod = tvm.build(tvm.IRModule.from_expr(tvm.te.create_prim_func([A])))
mod(*args)
return args[0].numpy()
diff --git a/python/tvm/tir/buffer.py b/python/tvm/tir/buffer.py
index 1109cc3d66..72c2a40fed 100644
--- a/python/tvm/tir/buffer.py
+++ b/python/tvm/tir/buffer.py
@@ -304,29 +304,6 @@ def decl_buffer(
buffer : tvm.tir.Buffer
The created buffer
- Example
- -------
- Here's an example of how broadcast buffer can be used to define a symbolic broadcast operation,
-
- .. code-block:: python
-
- m0, m1, m2 = te.var("m0"), te.var("m1"), te.var("m2")
- n0, n1, n2 = te.var("n0"), te.var("n1"), te.var("n2")
- o0, o1, o2 = te.var("o0"), te.var("o1"), te.var("o2")
- A = te.placeholder((m0, m1, m2), name='A')
- B = te.placeholder((n0, n1, n2), name='B')
- C = te.compute((o0, o1, o2), lambda i, j, k: A[i, j, k] + B[i, j, k], name='C')
- Ab = tvm.tir.decl_buffer(A.shape, A.dtype, name="Ab", buffer_type="auto_broadcast")
- Bb = tvm.tir.decl_buffer(B.shape, B.dtype, name="Bb", buffer_type="auto_broadcast")
- s = te.create_schedule(C.op)
- fadd = tvm.build(s, [A, B, C], target='llvm', name='bcast_add', binds={A:Ab, B:Bb})
- dev = tvm.cpu(0)
- a = tvm.nd.array(np.random.uniform(size=(2, 4, 3)).astype(A.dtype), dev)
- b = tvm.nd.array(np.random.uniform(size=(2, 1, 3)).astype(B.dtype), dev)
- c = tvm.nd.array(np.zeros((2, 4, 3), dtype=C.dtype), dev)
- fadd(a, b, c)
- tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
-
Note
----
Buffer data structure reflects the DLTensor structure in dlpack.
diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py
index 3588c04d8f..1de6941c99 100644
--- a/python/tvm/topi/__init__.py
+++ b/python/tvm/topi/__init__.py
@@ -39,9 +39,7 @@ from .broadcast import *
from .sort import *
from .scatter import *
from .scatter_elements import *
-from .sparse_fill_empty_rows import *
from .sparse_reshape import *
-from .argwhere import *
from .scan import *
from .einsum import *
from .unique import *
@@ -49,9 +47,7 @@ from .searchsorted import *
from .signal import *
from . import nn
from . import utils
-from . import vision
from . import image
-from . import random
from . import gpu
# error reporting
diff --git a/python/tvm/topi/argwhere.py b/python/tvm/topi/argwhere.py
deleted file mode 100644
index c2b658a4e9..0000000000
--- a/python/tvm/topi/argwhere.py
+++ /dev/null
@@ -1,197 +0,0 @@
-# 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.
-# pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks
-"""Argwhere operator"""
-import tvm
-from tvm.te import hybrid
-
-
-@hybrid.script
-def hybrid_argwhere_1d(output_shape, condition):
- """Find the indices of elements of a 1-D tensor that are non-zero.
-
- Parameters
- ----------
- condition : tvm.te.Tensor
- 1-D tensor with boolean values.
-
- Returns
- -------
- out : tvm.te.Tensor
- Indices of non-zero elements.
- """
- a = output_tensor(output_shape, "int32")
- a1 = condition.shape[0]
- valid_index = 0
- for i1 in range(a1):
- if condition[i1] != 0:
- a[valid_index, 0] = i1
- valid_index += 1
- return a
-
-
-@hybrid.script
-def hybrid_argwhere_2d(output_shape, condition):
- """Find the indices of elements of a 2-D tensor that are non-zero.
-
- Parameters
- ----------
- condition : tvm.te.Tensor
- 2-D tensor with boolean values.
-
- Returns
- -------
- out : tvm.te.Tensor
- Indices of non-zero elements.
- """
- a = output_tensor(output_shape, "int32")
- a1 = condition.shape[0]
- a2 = condition.shape[1]
- valid_index = 0
- for i1 in range(a1):
- for i2 in range(a2):
- if condition[i1, i2] != 0:
- a[valid_index, 0] = i1
- a[valid_index, 1] = i2
- valid_index += 1
- return a
-
-
-@hybrid.script
-def hybrid_argwhere_3d(output_shape, condition):
- """Find the indices of elements of a 3-D tensor that are non-zero.
-
- Parameters
- ----------
- condition : tvm.te.Tensor
- 3-D tensor with boolean values.
-
- Returns
- -------
- out : tvm.te.Tensor
- Indices of non-zero elements.
- """
- a = output_tensor(output_shape, "int32")
- a1 = condition.shape[0]
- a2 = condition.shape[1]
- a3 = condition.shape[2]
- valid_index = 0
- for i1 in range(a1):
- for i2 in range(a2):
- for i3 in range(a3):
- if condition[i1, i2, i3] != 0:
- a[valid_index, 0] = i1
- a[valid_index, 1] = i2
- a[valid_index, 2] = i3
- valid_index += 1
- return a
-
-
-@hybrid.script
-def hybrid_argwhere_4d(output_shape, condition):
- """Find the indices of elements of a 4-D tensor that are non-zero.
-
- Parameters
- ----------
- condition : tvm.te.Tensor
- 4-D tensor with boolean values.
-
- Returns
- -------
- out : tvm.te.Tensor
- Indices of non-zero elements.
- """
- a = output_tensor(output_shape, "int32")
- a1 = condition.shape[0]
- a2 = condition.shape[1]
- a3 = condition.shape[2]
- a4 = condition.shape[3]
- valid_index = 0
- for i1 in range(a1):
- for i2 in range(a2):
- for i3 in range(a3):
- for i4 in range(a4):
- if condition[i1, i2, i3, i4] != 0:
- a[valid_index, 0] = i1
- a[valid_index, 1] = i2
- a[valid_index, 2] = i3
- a[valid_index, 3] = i4
- valid_index += 1
- return a
-
-
-@hybrid.script
-def hybrid_argwhere_5d(output_shape, condition):
- """Find the indices of elements of a 5-D tensor that are non-zero.
-
- Parameters
- ----------
- condition : tvm.te.Tensor
- 5-D tensor with boolean values.
-
- Returns
- -------
- out : tvm.te.Tensor
- Indices of non-zero elements.
- """
- a = output_tensor(output_shape, "int32")
- a1 = condition.shape[0]
- a2 = condition.shape[1]
- a3 = condition.shape[2]
- a4 = condition.shape[3]
- a5 = condition.shape[4]
- valid_index = 0
- for i1 in range(a1):
- for i2 in range(a2):
- for i3 in range(a3):
- for i4 in range(a4):
- for i5 in range(a5):
- if condition[i1, i2, i3, i4, i5] != 0:
- a[valid_index, 0] = i1
- a[valid_index, 1] = i2
- a[valid_index, 2] = i3
- a[valid_index, 3] = i4
- a[valid_index, 4] = i5
- valid_index += 1
- return a
-
-
-@tvm.target.generic_func
-def argwhere(output_shape, condition):
- """Find the indices of elements of a tensor that are non-zero.
-
- Parameters
- ----------
- condition : tvm.te.Tensor
- Tensor with boolean values.
-
- Returns
- -------
- out : tvm.te.Tensor
- Indices of non-zero elements.
- """
- if len(condition.shape) == 1:
- return hybrid_argwhere_1d(output_shape.shape, condition)
- if len(condition.shape) == 2:
- return hybrid_argwhere_2d(output_shape.shape, condition)
- if len(condition.shape) == 3:
- return hybrid_argwhere_3d(output_shape.shape, condition)
- if len(condition.shape) == 4:
- return hybrid_argwhere_4d(output_shape.shape, condition)
- if len(condition.shape) == 5:
- return hybrid_argwhere_5d(output_shape.shape, condition)
- raise ValueError("Does not support rank higher than 5 in argwhere")
diff --git a/python/tvm/topi/nn/conv2d.py b/python/tvm/topi/nn/conv2d.py
index 5ee625577e..e145add5f0 100644
--- a/python/tvm/topi/nn/conv2d.py
+++ b/python/tvm/topi/nn/conv2d.py
@@ -615,68 +615,6 @@ def conv2d_NCHWc_int8(
)
-def conv2d_gemm_weight_transform(kernel, tile_N, tile_K, use_scalable_vectors=False, use_sme=False):
- """Weight transformation for winograd
-
- Parameters
- ----------
- kernel: Tensor
- The raw kernel tensor with layout "NHWC".
- tile_N: int
- Tile size across N axis of the weight transformation for ConvGemm. (N = OC)
- tile_K: int
- Tile size across K axis of the weight transformation for ConvGemm. (K = KW * KH * IC)
- use_scalable_vectors : bool
- determines if operations on scalable vectors are expected
- use_sme : bool
- determines if SME operations on scalable vectors are expected
-
- Returns
- -------
- output : tvm.te.Tensor
- 2-D with shape [CI*KH*KW,CO]
- """
- KH, KW, IC, OC = get_const_tuple(kernel.shape)
- K = KH * KW * IC
- N = OC
-
- kernel_flat = te.compute(
- (K, N), lambda x, y: kernel[(x // IC) // KW, (x // IC) % KW, x % IC, y], "weight_flatten"
- )
-
- pad_N, pad_K = tvm.topi.arm_cpu.arm_utils.get_conv2d_weights_padding(N, K, tile_N, tile_K)
-
- N_padded = N + pad_N
- K_padded = K + pad_K
-
- if pad_K != 0 or pad_N != 0:
- kernel_flat = pad(
- kernel_flat, pad_before=(0, 0), pad_after=(pad_K, pad_N), name="weight_padding"
- )
-
- if use_sme and kernel.dtype == "float16":
- return te.compute(
- (N_padded, K_padded), lambda x, y: kernel_flat[y, x], name="weight_transpose"
- )
-
- if use_scalable_vectors or use_sme:
- return kernel_flat
-
- if kernel.dtype in ["int8", "uint8"]:
- B_inter_t = te.compute(
- (N_padded // tile_N, K_padded // tile_K, tile_N, tile_K),
- lambda x, y, z, w: kernel_flat[w + tile_K * y, z + tile_N * x],
- name="weight_block_reshape",
- )
- else:
- B_inter_t = te.compute(
- (N_padded // tile_N, K_padded // tile_K, tile_K, tile_N),
- lambda x, y, z, w: kernel_flat[z + tile_K * y, w + tile_N * x],
- name="weight_block_reshape",
- )
- return B_inter_t
-
-
def conv2d_winograd_weight_transform(kernel, tile_size):
"""Weight transformation for winograd
@@ -712,29 +650,6 @@ def conv2d_winograd_weight_transform(kernel, tile_size):
)
-def conv2d_winograd_nnpack_weight_transform(kernel, convolution_algorithm, out_dtype):
- """Weight transformation for winograd
-
- Parameters
- ----------
- kernel: Tensor
- The raw kernel tensor with layout "NCHW". Only 3x3 kernel is supported for now.
- convolution_algorithm: int
- The convolution algorithm for Winograd NNPACK.
-
- Returns
- -------
- output : tvm.te.Tensor
- 4-D with shape [alpha, alpha, CO, CI]
- """
- # pylint: disable=import-outside-toplevel
- from tvm.contrib import nnpack
-
- return nnpack.convolution_inference_weight_transform(
- kernel, algorithm=convolution_algorithm, dtype=out_dtype
- )
-
-
def group_conv2d_nchw(Input, Filter, stride, padding, dilation, groups, out_dtype=None):
"""Group convolution operator in NCHW layout.
diff --git a/python/tvm/topi/random/__init__.py b/python/tvm/topi/random/__init__.py
deleted file mode 100644
index ee8d1d6385..0000000000
--- a/python/tvm/topi/random/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# 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.
-
-# pylint: disable=wildcard-import
-"""Pseudorandom generator kernels and operators."""
-from __future__ import absolute_import
-
-from .kernel import *
diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py
deleted file mode 100644
index 464ea9634a..0000000000
--- a/python/tvm/topi/random/kernel.py
+++ /dev/null
@@ -1,657 +0,0 @@
-# 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.
-"""Pseudorandom number kernels."""
-import math
-import numpy as np
-
-import tvm
-import tvm.topi
-
-from ... import tir
-from ...tir import ir_builder
-
-
-# Threefry PRNG with splitting based on
-# - J. K. Salmon, M. A. Moraes, R. O. Dror and D. E. Shaw, "Parallel random numbers: As easy as 1,
-# 2, 3," SC '11: Proceedings of 2011 International Conference for High Performance Computing,
-# Networking, Storage and Analysis, Seattle, WA, 2011, pp. 1-12, doi: 10.1145/2063384.2063405.
-# - Claessen, K. ; Palka, M. (2013) "Splittable Pseudorandom Number Generators using Cryptographic
-# Hashing". Proceedings of Haskell Symposium 2013 pp. 47-58. MLA
-# - Ferguson, Niels, et al. "The Skein hash function family." Submission to NIST (round 3) 7.7.5
-# (2010): 3.
-
-
-# Threefry is a counter based PRNG: given a unique input, it generates a unique random number. As
-# there is no state to maintain, we can apply it to a sequence of numbers (0..N) to generate a
-# sequence of random numbers in parallel. In order to make the PRNG splittable (that is we can
-# generate a sequence of random numbers in one place, and another sequence in another), we add a
-# path and key in addition to the counter. The path allows us to encode a sequence of splits (a 0 in
-# the path indicates the left result of a split, a 1 indicates the right). To avoid continuously
-# growing the path, we can compress an existing path into the key portion of the generator by
-# hashing the current key, path, and counter to create the new key (this same technique is used if
-# we run out of room for the counter). They key is initialized with a unique initial state.
-#
-# Random numbers are generated by applying the Threefry hash to the current key, path, and counter.
-
-# This module use encoding e4 from the appendix of "Splittable Pseudorandom Number Generators using
-# Cryptographic Hashing" (confusingly, the definition in the paper uses e3 to define the encoding
-# function). This encoding uses a 10 element uint64 tensor where each byte means the following:
-
-# .. code-block:
-
-# gen:
-# words: 0 1 2 3 | 4 5 | 6 7 | 8 9
-# usage: key | path | counter | position of next step in path encoded in binary
-# ex: 0b00010 -> next path entry goes one from the right
-
-# Right now, counter only uses the rightmost word.
-
-# Threefry rotation constants from the Skein paper ("The Skein Hash Function Family"
-# https://www.schneier.com/wp-content/uploads/2015/01/skein.pdf)
-_ROTATIONS = {
- 4: [[14, 16], [52, 57], [23, 40], [5, 37], [25, 33], [46, 12], [58, 22], [32, 32]],
- 8: [
- [46, 36, 19, 37],
- [33, 27, 14, 42],
- [17, 49, 36, 39],
- [44, 9, 54, 56],
- [39, 30, 34, 24],
- [13, 50, 10, 17],
- [25, 29, 39, 43],
- [8, 35, 56, 22],
- ],
- 16: [
- [24, 13, 8, 47, 8, 17, 22, 37],
- [38, 19, 10, 55, 49, 18, 23, 52],
- [33, 4, 51, 13, 34, 41, 59, 17],
- [5, 20, 48, 41, 47, 28, 16, 25],
- [41, 9, 37, 31, 12, 47, 44, 30],
- [16, 34, 56, 51, 4, 53, 42, 41],
- [31, 44, 47, 46, 19, 42, 44, 25],
- [9, 48, 35, 52, 23, 31, 37, 20],
- ],
-}
-
-# Threefry permutation constants from the Skein paper ("The Skein Hash Function Family"
-# https://www.schneier.com/wp-content/uploads/2015/01/skein.pdf)
-_PERMUTATIONS = {
- 4: [0, 3, 2, 1],
- 8: [2, 1, 4, 7, 6, 5, 0, 3],
- 16: [0, 9, 2, 13, 6, 11, 4, 15, 10, 7, 12, 3, 14, 5, 8, 1],
-}
-
-
-def _threefry(
- irb, key_buf, key_offset, counter_buf, counter_offset, out_buf, out_offset, out_shape
-):
- """IRBuilder code for running Threefry
-
- Parameters
- ----------
- irb: IRBuilder
- IRBuilder that this code will be generated for.
-
- key_buf: BufferVar
- Buffer to read the key from.
-
- key_offset: number
- Threefry will write to :code:`key_buf[key_offset:key_offset+4]`
-
- counter_buf: BufferVar
- Buffer to read the counter from.
-
- counter_offset: number
- Threefry will write to :code:`counter_buf[counter_offset:counter_offset+4]`
-
- out_buf: BufferVar
- Buffer to read the counter from.
-
- out_offset: number
- Threefry will write to :code:`out_buf[out_offset:out_offset+4*product(out_shape)]`
-
- out_shape: number
- Determines the number of output states to generate. :code:`state[i]` will correspond to
- counter+i.
- """
- nrounds = 20
- nwords = 4
- iwidth = 64
- assert nrounds % 4 == 0
- assert nwords in [4, 8, 16]
-
- # The paper has constants for 32 bit threefry, but we keep the implementation simple by only
- # using 64-bit words.
- assert key_buf.dtype == "uint64", "threefry only supports 64-bit keys"
- assert key_buf.dtype == counter_buf.dtype, "threefry key and counter must be the same dtype"
-
- def mix(a, b, rotation):
- x = a + b # wrapping
- y = x ^ ((b << rotation) | (b >> (iwidth - rotation)))
- return [x, y]
-
- # temporary buffer for holding the results of _PERMUTATIONS
- tmp = irb.allocate(out_buf.dtype, out_shape * nwords, name="tmp", scope="global")
- tmp_offset = 0
-
- # Initialize entire key. It is composed of the original key with one
- # element appended. The appended element is the xor of all key words plus a
- # constant.
- full_key = irb.allocate("uint64", nwords + 1, name="full_key", scope="global")
- for i in range(nwords):
- full_key[i] = key_buf[key_offset + i]
- # initial key constant, full_key[nwords] is equivalent to k_{N_W} in the Skein paper.
- full_key[nwords] = tvm.tir.const(0x1BD11BDAA9FC1A22, dtype="uint64")
- for i in range(nwords):
- full_key[nwords] ^= key_buf[key_offset + i]
-
- with irb.for_range(0, out_shape, dtype="uint64", name="i") as i:
- for j in range(nwords):
- out_buf[out_offset + i * nwords + j] = counter_buf[counter_offset + j] + i
-
- def key_schedule(s, i):
- # Threefry uses no tweak, so the key schedule is simple
- if i == nwords - 1:
- return full_key[(s + i) % (nwords + 1)] + tvm.tir.const(s, dtype="uint64")
- return full_key[(s + i) % (nwords + 1)]
-
- with irb.for_range(0, out_shape, name="l") as l: # pylint: disable=invalid-name
- for i in range(nrounds // 4):
- for j in range(nwords):
- out_buf[out_offset + l * nwords + j] += key_schedule(i, j) # wrapping
- for k in range(4):
- for j in range(nwords // 2):
- (
- out_buf[out_offset + l * nwords + j * 2 + 0],
- out_buf[out_offset + l * nwords + j * 2 + 1],
- ) = mix(
- out_buf[out_offset + l * nwords + j * 2 + 0],
- out_buf[out_offset + l * nwords + j * 2 + 1],
- _ROTATIONS[nwords][(i * 4 + k) % 8][j],
- )
- for j in range(nwords):
- tmp[tmp_offset + l * nwords + j] = out_buf[
- out_offset + l * nwords + _PERMUTATIONS[nwords][j]
- ]
- # number of rounds is even, so out always contains the result
- (out_buf, tmp) = (tmp, out_buf)
- (out_offset, tmp_offset) = (tmp_offset, out_offset)
-
-
-def threefry_generate(gen, out_shape):
- """Generate a series of random values
-
- Notes
- -----
- This function uses the counter portion of the generator state to generate a series of random
- numbers in parallel. Random number `i` is generated by applying Threefry to the current
- generator state with the counter portion incremented by `i`. This means that each random number
- is generated independently from each other random number, so we can compute them in parallel.
-
- If there is not enough room left in the counter to generate the desired shape of random values,
- then a new generator is created by applying Threefry to the current key, path, and counter.
- This new generator will have a reset counter.
-
- Warning
- -------
- Threeyfry requires that unsigned integer arithmetic wraps on overflow. Currently TVM has no
- guarantee of this, so threefry contains an internal assert to check wrapping behavior. This
- assert may or may not run depending on your platform, so it is recommended you run
- :py:func:`threefry_test_wrapping` to verify wrapping behavior.
-
- Parameters
- ----------
- gen : Tensor[10, uint64]
- Generator state. Can be create with :py:func:`tvm.relay.random.threefry_key`. This should
- not be reused in another function, otherwise random numbers will be repeated.
-
- out_shape : Sequence[int]
- Output shape of the random numbers.
-
- Returns
- -------
- new_gen : Tensor[10, uint64]
- The new generator state to be used in subsequent calls.
-
- rand : Tensor[out_shape, uint64]
- Tensor of random numbers with shape `out_shape`.
- """
- out_len = tir.const(1)
- for s in out_shape:
- out_len *= s
- assert (
- out_len.value <= 2**64 - 1
- ), f"Can only generate up to 2^64 random numbers, but {out_len} were requested."
-
- def gen_ir(gen_ptr, out_gen_ptr, out_array_ptr):
- irb = ir_builder.create()
- gen = irb.buffer_ptr(gen_ptr)
- out_gen = irb.buffer_ptr(out_gen_ptr)
- out_array = irb.buffer_ptr(out_array_ptr)
-
- # Check that unsigned arithmetic wraps, as it is required to implement threefry correctly.
- irb.emit(
- tvm.tir.AssertStmt(
- tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64") + tvm.tir.const(1, "uint64")
- == tvm.tir.const(0, "uint64"),
- tvm.tir.StringImm(
- "Unsigned integer arithmetic is not wrapping, but threefry requires wrapping."
- ),
- tvm.tir.Evaluate(0),
- )
- )
-
- # Create a temporary array to hold the generator state we will use to create the random
- # numbers. We cannot use gen because we may need to update the key + path if there is not
- # enough room in the counter.
- tmp = irb.allocate(gen.dtype, 10, name="tmp", scope="global")
-
- # TODO(tkonolige): for now we only use the last word of the counter for counting. It is too
- # much work to figure out how to do 128 bit addition.
-
- # Max value for counter should be 2**64-2 because we need to reserve a special value to
- # indicate the counter is used up.
- with irb.if_scope(gen[7] < tir.const(2**64 - 1, dtype=gen.dtype) - out_len):
- for i in range(10):
- tmp[i] = gen[i]
- with irb.else_scope():
- # no room left in the counter, we have to change the path or key
- with irb.if_scope(gen[8] == 0 and gen[9] == 0):
- # out of room in the path, have to generate new key
-
- # The paper says the counter that we will be hashing should be a special value of
- # all ones. We need to allocate some space for it because we cannot overwrite gen.
- tmp_counter = irb.allocate(gen.dtype, 2, name="tmp_counter", scope="global")
- tmp_counter[0] = tir.const(0xFFFFFFFFFFFFFFFF, dtype=gen.dtype)
- tmp_counter[1] = tir.const(0xFFFFFFFFFFFFFFFF, dtype=gen.dtype)
- _threefry(irb, gen, 0, tmp_counter, 0, tmp, 0, 1)
- tmp[4] = tir.const(0, dtype=gen.dtype) # zero path, i.e. no path
- tmp[5] = tir.const(0, dtype=gen.dtype)
- tmp[6] = tir.const(0, dtype=gen.dtype) # zero counter
- tmp[7] = tir.const(0, dtype=gen.dtype)
- tmp[8] = tir.const(1 << 63, dtype=gen.dtype) # one in the leftmost position
- tmp[9] = tir.const(0, dtype=gen.dtype)
- with irb.else_scope():
- tmp[0] = gen[0]
- tmp[1] = gen[1]
- tmp[2] = gen[2]
- tmp[3] = gen[3]
- tmp[4] = gen[4] | gen[8] # add a 1 to the path
- tmp[5] = gen[5] | gen[9]
- tmp[6] = tir.const(0, dtype=gen.dtype) # zero counter
- tmp[7] = tir.const(0, dtype=gen.dtype)
- _shift_right(irb, gen[8], gen[9], tmp, 8, tmp, 9)
-
- # Compute random values
- if out_len.value >= 4:
- _threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4)
- if out_len.value % 4 != 0:
- remaining = irb.allocate(gen.dtype, 4, name="remaining", scope="global")
- tmp[7] = tmp[7] + tir.Cast(gen.dtype, out_len // 4 * 4) # increment counter
- _threefry(irb, tmp, 0, tmp, 4, remaining, 0, 1)
- with irb.for_range(0, out_len % 4, dtype="uint64", name="i") as i:
- out_array[out_len // 4 * 4 + i] = remaining[i]
-
- # Update generator state
- out_gen[0] = tmp[0] # key stays the same
- out_gen[1] = tmp[1]
- out_gen[2] = tmp[2]
- out_gen[3] = tmp[3]
- out_gen[4] = tmp[4] # path stays the same
- out_gen[5] = tmp[5]
- out_gen[6] = tir.const(0, dtype=gen.dtype) # unused, leave it as 0
- if out_len.value % 4 != 0:
- # increment counter for the remaining
- # as we will generate 4 random numbers for the remaining, increase 4 here.
- # the main increment was done before the second _threefry.
- out_gen[7] = tmp[7] + tir.Cast(gen.dtype, 4)
- else:
- out_gen[7] = tmp[7] + tir.Cast(gen.dtype, out_len) # increment counter
- out_gen[8] = tmp[8] # path unchanged, so no update here
- out_gen[9] = tmp[9]
-
- return irb.get()
-
- out_gen = tvm.tir.decl_buffer((10,), name="out_gen", dtype="uint64")
- out_array = tvm.tir.decl_buffer(out_shape, name="out_array", dtype="uint64")
- return tvm.te.extern(
- [out_gen.shape, out_array.shape],
- [gen],
- lambda ins, outs: gen_ir(ins[0], outs[0], outs[1]),
- out_buffers=[out_gen, out_array],
- name="threefry_generate",
- tag="threefry_generate",
- )
-
-
-def _shift_right(irb, a, b, out_a, a_off, out_b, b_off):
- """Binary shift a 128bit number composed of two 64 bit words right by one."""
- with irb.if_scope(a == 1):
- out_a[a_off] = tir.const(0, dtype=a.dtype)
- out_b[b_off] = tir.const(0x8000000000000000, dtype=a.dtype)
- with irb.else_scope():
- with irb.if_scope(a == 0):
- out_a[a_off] = tir.const(0, dtype=a.dtype)
- out_b[b_off] = b >> 1
- with irb.else_scope():
- out_a[a_off] = a >> 1
- out_b[b_off] = tir.const(0, dtype=a.dtype)
-
-
-def threefry_split(gen):
- """Split a single generator state into two new ones
-
- Notes
- -----
- The new generator is created by appending a one (for the right output) or a zero (for the left
- output) to the end of the path portion of the generator If there is no longer and room in the
- path, then we create a new key portion of the generator by applying Threefry to the old state,
- path, and counter. i.e. :code:`new_key = threefry(old_key, [old_path, old_counter])`. This
- resets the path portion of the new generator.
-
- Parameters
- ----------
- gen : Tensor[10, uint64]
- Generator state. Can be create with :py:func:`tvm.relay.random.threefry_key`. This should
- not be reused in another function, otherwise random numbers will be repeated.
-
- Returns
- -------
- out_gen_left : Tensor[10, uint64]
- New generator state that is distinct from `out_gen_right`.
-
- out_gen_right : Tensor[10, uint64]
- New generator state that is distinct from `out_gen_left`.
- """
-
- def gen_ir(gen_ptr, out_left_ptr, out_right_ptr):
- irb = ir_builder.create()
- gen = irb.buffer_ptr(gen_ptr)
- out_left = irb.buffer_ptr(out_left_ptr)
- out_right = irb.buffer_ptr(out_right_ptr)
-
- with irb.if_scope(gen[8] == 0 and gen[9] == 0):
- # Generate new key because we have run out of room to extend the path
- _threefry(irb, gen, 0, gen, 4, out_left, 0, 1)
- out_left[4] = tir.const(0, dtype=gen.dtype)
- out_left[5] = tir.const(0, dtype=gen.dtype)
- out_left[6] = tir.const(0, dtype=gen.dtype) # counter gets zeroed
- out_left[7] = tir.const(0, dtype=gen.dtype) # counter gets zeroed
- out_left[8] = tir.const(
- 1 << 62, dtype=gen.dtype
- ) # one in the second from the leftmost position
- out_left[9] = tir.const(0, dtype=gen.dtype)
-
- out_right[0] = out_left[0]
- out_right[1] = out_left[1]
- out_right[2] = out_left[2]
- out_right[3] = out_left[3]
- out_right[4] = tir.const(1 << 63, dtype=gen.dtype) # one in the leftmost position
- out_right[5] = tir.const(0, dtype=gen.dtype)
- out_right[6] = tir.const(0, dtype=gen.dtype)
- out_right[7] = tir.const(0, dtype=gen.dtype)
- out_right[8] = tir.const(
- 1 << 62, dtype=gen.dtype
- ) # one in the second from the leftmost position
- out_right[9] = tir.const(0, dtype=gen.dtype)
- with irb.else_scope():
- out_left[0] = gen[0]
- out_left[1] = gen[1]
- out_left[2] = gen[2]
- out_left[3] = gen[3]
- out_left[4] = gen[4] # adding a zero here, but its already zero padded
- out_left[5] = gen[5]
- out_left[6] = gen[6]
- out_left[7] = gen[7]
- # move path position over one bit
- _shift_right(irb, gen[8], gen[9], out_left, 8, out_left, 9)
-
- out_right[0] = gen[0]
- out_right[1] = gen[1]
- out_right[2] = gen[2]
- out_right[3] = gen[3]
- out_right[4] = gen[4] | gen[8] # add a one to the path
- out_right[5] = gen[5] | gen[9]
- out_right[6] = gen[6]
- out_right[7] = gen[7]
- _shift_right(irb, gen[8], gen[9], out_right, 8, out_right, 9)
-
- return irb.get()
-
- out_left = tvm.tir.decl_buffer((10,), name="out_left", dtype="uint64")
- out_right = tvm.tir.decl_buffer((10,), name="out_right", dtype="uint64")
- return tvm.te.extern(
- [out_left.shape, out_right.shape],
- [gen],
- lambda ins, outs: gen_ir(ins[0], outs[0], outs[1]),
- out_buffers=[out_left, out_right],
- name="threefry_split",
- tag="threefry_split",
- )
-
-
-def threefry_test_wrapping(target, device):
- """Test that unsigned arithmetic wraps on overflow.
-
- Parameters
- ----------
- target : tvm.target.Target
- Target to run against
- device : tvm.runtime.Device
- Context to run the test on
-
- Returns
- -------
- is_wrapping : bool
- Whether or not unsigned integer arithmetic is wrapping for this target, context pair. True
- indicates that threefry will work on this platform.
- """
- if isinstance(target, str):
- target = tvm.target.Target(target)
-
- def gen_ir(out_ptr):
- irb = ir_builder.create()
- out = irb.buffer_ptr(out_ptr)
- if "gpu" in target.keys:
- thread_x = tvm.te.thread_axis("threadIdx.x")
- irb.scope_attr(thread_x, "thread_extent", 1)
- out[0] = tvm.tir.const(0xFFFFFFFFFFFFFFFF, "uint64") + tvm.tir.const(1, "uint64")
- return irb.get()
-
- out = tvm.tir.decl_buffer((1,), dtype="uint64")
- f = tvm.te.extern(
- [out.shape], [], lambda ins, outs: gen_ir(outs[0]), dtype="uint64", out_buffers=[out]
- )
- s = tvm.te.create_schedule([f.op])
- out_ary = tvm.nd.array(np.ones((1,), "uint64"), device)
- tvm.build(s, [f], target=target)(out_ary)
- return out_ary.numpy()[0] == 0
-
-
-def uniform(gen, low, high, out_shape, out_dtype):
- """Draw samples from a uniform distribution.
-
- Samples are uniformly distributed over the half-open interval [low, high)
- (includes low, but excludes high). In other words, any value within the
- given interval is equally likely to be drawn by uniform.
-
- Parameters
- ----------
- gen : ThreefryKey
- Generator state. Can be create with :py:func:`tvm.relay.threefry_key`. This should not be
- reused in another function, otherwise random numbers will be repeated.
-
- low : Tensor[(), out_dtype]
- Lower boundary of the output interval. All values generated will be
- greater than or equal to low.
-
- high : Tensor[(), out_dtype]
- Upper boundary of the output interval. All values generated will be
- less than high.
-
- out_shape : Sequence[int]
- Output shape of the random numbers.
-
- out_dtype : str
- The output dtype.
-
- Returns
- -------
- new_gen : ThreefryKey
- New generator state that is distinct from `gen`.
-
- out : Tensor[out_shape, out_dtype]
- Tensor of random numbers with shape `out_shape` and type `out_dtype`.
- """
- new_gen, random_bits = threefry_generate(gen, out_shape)
- assert out_dtype in (
- "float32",
- "float64",
- ), f"Only support float32 or float64 for now, got {out_dtype}"
- if out_dtype == "float32":
- random_dtype = "uint32"
- nbits = 32
- nfraction = 23
- elif out_dtype == "float64":
- random_dtype = "uint64"
- nbits = 64
- nfraction = 52
- nexp = nbits - nfraction - 1
- random_bits = random_bits.astype(random_dtype)
-
- fraction = tvm.topi.right_shift(
- random_bits, tvm.tir.const(nbits - nfraction, dtype=random_dtype)
- )
- exponent = tvm.topi.left_shift(
- tvm.topi.full(out_shape, random_dtype, (1 << (nexp - 1)) - 1),
- tvm.tir.const(nfraction, dtype=random_dtype),
- )
- mantissa = tvm.topi.bitwise_or(fraction, exponent).astype(random_dtype)
- standard_uniform_values = tvm.topi.reinterpret(mantissa, out_dtype) - tvm.tir.const(
- 1, dtype=out_dtype
- )
- uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low)
-
- return new_gen, uniform_values
-
-
-def normal(gen, mean, scale, out_shape, out_dtype):
- """Draw samples from a normal distribution.
- The algorithm is based on Box-Muller transform
-
- Parameters
- ----------
- gen : ThreefryKey
- Generator state. Can be create with :py:func:`tvm.relay.threefry_key`. This should not be
- reused in another function, otherwise random numbers will be repeated.
-
- mean : Tensor[(), out_dtype]
- The mean of the normal distribution.
-
- scale : Tensor[(), out_dtype]
- The standard deviation of the normal distribution.
-
- out_shape : Sequence[int]
- Output shape of the random numbers.
-
- out_dtype : str
- The output dtype.
-
- Returns
- -------
- new_gen : ThreefryKey
- New generator state that is distinct from `gen`.
-
- out : Tensor[out_shape, out_dtype]
- Tensor of random numbers with shape `out_shape` and type `out_dtype`.
- """
- out_shape = list(out_shape)
- # Box-Muller transform need two pieces of original uniform data
- out_shape.insert(0, 2)
- new_gen, uniform_values = uniform(
- gen, tvm.tir.const(0.0, out_dtype), tvm.tir.const(1.0, out_dtype), out_shape, out_dtype
- )
- two_pi = tvm.tir.const(2.0 * math.pi, out_dtype)
- uniform_values_1 = tvm.topi.strided_slice(uniform_values, [0], [1], strides=[1], axes=[0])
- uniform_values_1 = tvm.topi.squeeze(uniform_values_1, axis=0)
- uniform_values_2 = tvm.topi.strided_slice(uniform_values, [1], [2], strides=[1], axes=[0])
- uniform_values_2 = tvm.topi.squeeze(uniform_values_2, axis=0)
- uniform_values_1 = tvm.topi.subtract(tvm.tir.const(1.0, out_dtype), uniform_values_1)
- sqrt_values = tvm.topi.sqrt(
- tvm.topi.multiply(tvm.tir.const(-2.0, out_dtype), tvm.topi.log(uniform_values_1))
- )
- sin_values = tvm.topi.sin(tvm.topi.multiply(two_pi, uniform_values_2))
- random_values = tvm.topi.add(
- tvm.topi.multiply(tvm.topi.multiply(sqrt_values, sin_values), scale), mean
- )
-
- return new_gen, random_values
-
-
-def multinomial(gen, probs, num_samples):
- """Draw samples from a multinomial distribution defined by the input tensor.
-
- Parameters
- ----------
- gen : ThreefryKey
- Generator state. Can be created with :py:func:`tvm.relay.threefry_key`. This should not be
- reused in another function, otherwise random numbers will be repeated.
-
- probs: Tensor[(input_rows, indices), float]
- A tensor containing the probabilities to sample from. Each value represents the
- probability of choosing its corresponding index. If a tensor is provided, the last dimension
- is treated independently. Negative values in this tensor will be clipped to zero to
- represent they have no chance of being selected.
-
- num_samples: int
- Number of samples to draw from each row.
-
- Returns
- -------
- new_gen : ThreefryKey
- New generator state that is distinct from `gen`.
-
- out : Tensor[(input_rows, num_samples), int64]
- Tensor of sampled indices with shape `input_rows x num_samples` and type `out_dtype`.
- """
- # Convert to float for consistent behavior.
- probs = tvm.topi.cast(probs, "float32")
- # Clip negative values to 0.
- probs = tvm.topi.maximum(probs, 0)
- # Normalize input probabilities.
- probs = tvm.topi.divide(probs, tvm.topi.expand_dims(tvm.topi.sum(probs, axis=-1), -1))
- # Convert probability to cumulative sum.
- cumulative_probs = tvm.topi.cumsum(probs, axis=-1)
- # Sample a set of uniform values.
- new_gen, uniform_values = uniform(
- gen,
- tvm.tir.const(0.0, "float32"),
- tvm.tir.const(1.0, "float32"),
- [*probs.shape[:-1], num_samples],
- "float32",
- )
- # Find index corresponding to sampled values.
- closest_prob = tvm.topi.subtract(
- tvm.topi.expand_dims(cumulative_probs, axis=-1),
- tvm.topi.expand_dims(uniform_values, axis=-2),
- )
- zeros = tvm.topi.full_like(closest_prob, 0)
- ones = tvm.topi.full_like(closest_prob, 1)
- # Find the smallest positive index for each sample.
- cond = tvm.topi.greater(closest_prob, zeros)
- closest_non_neg = tvm.topi.where(cond, closest_prob, ones)
- sampled_indices = tvm.topi.argmin(closest_non_neg, axis=-2)
- return new_gen, sampled_indices
diff --git a/python/tvm/topi/sparse_fill_empty_rows.py b/python/tvm/topi/sparse_fill_empty_rows.py
deleted file mode 100644
index 10dc6ee3bf..0000000000
--- a/python/tvm/topi/sparse_fill_empty_rows.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# 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, WITHnew_sparse_indices WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-# pylint: disable=no-else-return, too-many-locals, too-many-arguments, too-many-branches
-# pylint: disable=undefined-variable, invalid-name
-"""SparseFillEmptyRows operator"""
-from ..te import hybrid
-
-
-@hybrid.script
-def _sparse_fill_empty_rows(
- sparse_indices,
- sparse_values,
- dense_shape,
- default_value,
- new_sparse_indices_shape,
- new_sparse_values_shape,
- empty_row_indicator_shape,
-):
- default_value_ = int64(default_value[0])
- new_sparse_indices = output_tensor(new_sparse_indices_shape, "int64")
- new_sparse_values = output_tensor(new_sparse_values_shape, "int64")
- empty_row_indicator = output_tensor(empty_row_indicator_shape, "int64")
- new_sparse_indices_row_id = 0
-
- if int64(sparse_indices.shape[0]) == int64(0): # Handle Empty Case
- # Fill all rows with default values
- for i in range(0, new_sparse_indices_shape[0]):
- new_sparse_indices[i, 0] = int64(i)
- new_sparse_values[i] = default_value_
- empty_row_indicator[i] = int64(1)
- for k in range(1, int64(new_sparse_indices_shape[1])):
- new_sparse_indices[i, k] = int64(0)
-
- return (new_sparse_indices, new_sparse_values, empty_row_indicator)
-
- else:
- # Iterate through sparse_indices and add rows if/when required
- for i in range(0, int64(sparse_indices.shape[0])):
- if i == 0:
- prev_row_id = int64(0)
- else:
- prev_row_id = int64(sparse_indices[i - 1, 0] + 1)
- row_id = int64(sparse_indices[i, 0])
-
- # Since input is in row-major order, add rows between prev_row_id and row_id
- for j in range(prev_row_id, row_id):
- new_sparse_indices[new_sparse_indices_row_id, 0] = int64(j)
- for k in range(1, int64(new_sparse_indices_shape[1])):
- new_sparse_indices[new_sparse_indices_row_id, k] = int64(0)
- empty_row_indicator[prev_row_id] = int64(1)
- new_sparse_values[new_sparse_indices_row_id] = default_value_
- new_sparse_indices_row_id += 1
-
- # Add current element to output
- new_sparse_indices[new_sparse_indices_row_id, 0] = row_id
- for k in range(1, int64(new_sparse_indices_shape[1])):
- new_sparse_indices[new_sparse_indices_row_id, k] = int64(sparse_indices[i, k])
- new_sparse_values[new_sparse_indices_row_id] = int64(sparse_values[i])
- empty_row_indicator[row_id] = int64(0)
- new_sparse_indices_row_id += 1
-
- # Add rows with default value if last row id of sparse_indices is not dense_shape[0] - 1
- for i in range(
- int64(sparse_indices[sparse_indices.shape[0] - 1, 0] + 1), int64(dense_shape[0])
- ):
-
- new_sparse_indices[new_sparse_indices_row_id, 0] = int64(i)
- for k in range(1, int64(new_sparse_indices_shape[1])):
- new_sparse_indices[new_sparse_indices_row_id, k] = int64(0)
- empty_row_indicator[i] = int64(1)
- new_sparse_values[new_sparse_indices_row_id] = default_value_
- new_sparse_indices_row_id += 1
-
- return (new_sparse_indices, new_sparse_values, empty_row_indicator)
-
-
-def sparse_fill_empty_rows(
- sparse_indices,
- sparse_values,
- dense_shape,
- default_value,
- new_sparse_indices_shape,
- new_sparse_values_shape,
- empty_row_indicator_shape,
-):
- return _sparse_fill_empty_rows(
- sparse_indices,
- sparse_values,
- dense_shape,
- default_value,
- new_sparse_indices_shape,
- new_sparse_values_shape,
- empty_row_indicator_shape,
- )
diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 2844825a4a..c1f5bce948 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -20,7 +20,6 @@ from __future__ import absolute_import as _abs
import tvm
from tvm import te, topi
-from tvm.te import hybrid
from . import cpp, tag
from .utils import const_vector, make_idx, within_index
@@ -982,35 +981,6 @@ def adv_index(data, indices):
return cpp.adv_index(data, indices)
-@hybrid.script
-def invert_permutation(data):
- """Computes the inverse permutation of data.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- Input data
-
- Returns
- -------
- result : tvm.te.Tensor
- Output tensor
-
- Examples
- --------
- .. code-block:: python
-
- data = [3, 4, 0, 2, 1]
- topi.invert_permutation(data) = [2, 4, 3, 0, 1]
- """
- result = output_tensor(data.shape, data.dtype)
- nums = data.shape[0]
- for ind in range(nums):
- r_ind = data[ind]
- result[r_ind] = ind
- return result
-
-
def sliding_window(data, axis, window_shape, strides):
"""Slide a window over the data tensor.
diff --git a/python/tvm/topi/unique.py b/python/tvm/topi/unique.py
index 983c486153..9c97320134 100644
--- a/python/tvm/topi/unique.py
+++ b/python/tvm/topi/unique.py
@@ -17,9 +17,6 @@
# pylint: disable=invalid-name
"""Unique operator"""
from tvm import te, tir
-from ..te import hybrid
-from .scan import cumsum
-from .sort import sort, argsort
def _calc_adjacent_diff_ir(data, output, binop=tir.Sub):
@@ -82,234 +79,3 @@ def _calc_adjacent_diff(data, out_dtype="int32", binop=tir.Sub):
name="_calc_adjacent_diff",
tag="_calc_adjacent_diff_cpu",
)
-
-
-@hybrid.script
-def _calc_num_unique(inc_scan):
- """Helper function to get the number of unique elements fron inc_scan tensor"""
- output = output_tensor((1,), "int32")
- output[0] = inc_scan[inc_scan.shape[0] - 1] + int32(1)
- return output
-
-
-def _calc_unique_ir(
- data, argsorted_indices, inc_scan, index_converter, unique_elements, inverse_indices, counts
-):
- """Low level IR to calculate unique elements, inverse indices, and counts (optional) of
- unique elements of 1-D array.
-
- Parameters
- ----------
- data : Buffer
- Input 1-D Buffer.
-
- argsorted_indices : Buffer
- A buffer that stores the argsorted indices of the input data.
-
- inc_scan : Buffer
- A buffer that stores the inclusive scan of the binary tir.NE adjacent difference
- of the sorted data.
-
- index_converter (optional) : Buffer
- An optional index converter that transforms the unique element index
- such that new_idx = index_converter[old_idx].
-
- unique_elements : Buffer
- A buffer that stores the unique elements.
-
- inverse_indices : Buffer
- A buffer that stores the index of each input data element in the unique element array.
-
- counts (optional) : Buffer
- A buffer that stores the count of each unique element.
- """
- ib = tir.ir_builder.create()
- data_ptr = ib.buffer_ptr(data)
- argsorted_indices_ptr = ib.buffer_ptr(argsorted_indices)
- inc_scan_ptr = ib.buffer_ptr(inc_scan)
- unique_elements_ptr = ib.buffer_ptr(unique_elements)
- inverse_indices_ptr = ib.buffer_ptr(inverse_indices)
-
- index_converter_ptr = None
- if isinstance(index_converter, tir.Buffer):
- index_converter_ptr = ib.buffer_ptr(index_converter)
-
- if isinstance(counts, tir.Buffer):
- counts_ptr = ib.buffer_ptr(counts)
- # use indices_ptr as a tmp buffer to store tids with inc_scan[tid] != inc_scan[tid-1]
- unique_seq_indices_ptr = ib.buffer_ptr(inverse_indices)
-
- data_length = data.shape[0]
-
- # if need to return counts
- if isinstance(counts, tir.Buffer):
- num_unique = inc_scan_ptr[inc_scan.shape[0] - 1] + 1
- num_elements = data.shape[0]
- unique_seq_indices_ptr[num_unique - 1] = num_elements
- with ib.new_scope():
- with ib.for_range(0, data_length, kind="parallel") as i:
- with ib.if_scope(i > 0):
- with ib.if_scope(inc_scan_ptr[i] != inc_scan_ptr[i - 1]):
- unique_seq_indices_ptr[inc_scan_ptr[i] - 1] = i
- with ib.new_scope():
- with ib.for_range(0, num_unique, kind="parallel") as i:
- unique_idx = i if not index_converter_ptr else index_converter_ptr[i]
- with ib.if_scope(i == 0):
- counts_ptr[unique_idx] = unique_seq_indices_ptr[i]
- with ib.else_scope():
- counts_ptr[unique_idx] = (
- unique_seq_indices_ptr[i] - unique_seq_indices_ptr[i - 1]
- )
- # calculate unique elements and inverse indices
- with ib.new_scope():
- with ib.for_range(0, data_length, kind="parallel") as i:
- data_idx = argsorted_indices_ptr[i]
- unique_idx = (
- inc_scan_ptr[i] if not index_converter_ptr else index_converter_ptr[inc_scan_ptr[i]]
- )
- inverse_indices_ptr[data_idx] = unique_idx
- with ib.if_scope(i == 0):
- unique_elements_ptr[unique_idx] = data_ptr[data_idx]
- with ib.else_scope():
- with ib.if_scope(inc_scan_ptr[i] != inc_scan_ptr[i - 1]):
- unique_elements_ptr[unique_idx] = data_ptr[data_idx]
- return ib.get()
-
-
-@hybrid.script
-def _calc_first_occurence(argsorted_indices, inc_scan):
- """Hybrid script to calculate the first occurence of each unique element in the input data.
-
- Parameters
- ----------
- argsorted_indices : tvm.te.Tensor
- A tensor that stores the argsorted indices of the input data.
-
- inc_scan : tvm.te.Tensor
- A tensor that stores the inclusive scan of the binary tir.NE adjacent difference
- of the sorted data.
-
- first_occurence : tvm.te.Tensor
- A tensor that stores the first occurence of each unique element in the input data.
- """
- first_occurence = output_tensor(argsorted_indices.shape, "int32")
- for i in parallel(argsorted_indices.shape[0]):
- first_occurence[i] = argsorted_indices.shape[0]
- for i in parallel(argsorted_indices.shape[0]):
- if i == 0 or inc_scan[i] != inc_scan[i - 1]:
- first_occurence[inc_scan[i]] = argsorted_indices[i]
- return first_occurence
-
-
-def unique(data, is_sorted=True, return_counts=False):
- """
- Find the unique elements of a 1-D tensor. Please note `output` and `counts` are all padded to
- have the same length of `data` and element with index >= num_unique[0] has undefined value.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- A 1-D tensor of integers.
-
- sorted : bool
- Whether to sort the unique elements in ascending order before returning as output.
-
- return_counts : bool
- Whether to return the count of each unique element.
-
- Returns
- -------
- unique : tvm.te.Tensor
- A 1-D tensor containing the unique elements of the input data tensor. The same size as
- the input data. If there are less unique elements than input data, the end of the tensor
- is padded with zeros.
-
- indices : tvm.te.Tensor
- A 1-D tensor. The same size as output. For each entry in output, it contains
- the index of its first occurence in the input data. The end of the tensor is padded
- with the length of the input data.
-
- inverse_indices : tvm.te.Tensor
- A 1-D tensor. For each entry in data, it contains the index of that data element in
- the unique array. (Note that inverse_indices is very similar to indices if output is not
- sorted.)
-
- num_unique : tvm.te.Tensor
- A 1-D tensor with size=1 containing the number of unique elements in the input data tensor.
-
- counts (optional) : tvm.te.Tensor
- A 1-D tensor containing the count of each unique element in the output.
-
- Examples
- --------
- .. code-block:: python
-
- [output, indices, num_unique] = unique([4, 5, 1, 2, 3, 3, 4, 5], False, False)
- output = [4, 5, 1, 2, 3, _, _, _]
- indices = [0, 1, 2, 3, 4, _, _, _]
- inverse_indices = [0, 1, 2, 3, 4, 4, 0, 1]
- num_unique = [5]
-
- [output, indices, num_unique, counts] = unique([4, 5, 1, 2, 3, 3, 4, 5], False, True)
- output = [4, 5, 1, 2, 3, _, _, _]
- indices = [0, 1, 2, 3, 4, _, _, _]
- inverse_indices = [0, 1, 2, 3, 4, 4, 0, 1]
- num_unique = [5]
- counts = [2, 2, 1, 1, 2, _, _, _]
-
- [output, indices, num_unique] = unique([4, 5, 1, 2, 3, 3, 4, 5], True)
- output = [1, 2, 3, 4, 5, _, _, _]
- indices = [2, 3, 4, 0, 1, _, _, _]
- inverse_indices = [3, 4, 0, 1, 2, 2, 3, 4]
- num_unique = [5]
- """
- sorted_data = sort(data)
- argsorted_indices = argsort(data, dtype="int32")
- # adjacent difference
- adjacent_diff = _calc_adjacent_diff(sorted_data, "int32", tir.NE)
- # inclusive scan
- inc_scan = cumsum(adjacent_diff, dtype="int32", exclusive=0)
- # total number of unique elements
- num_unique_elements = _calc_num_unique(inc_scan)
- # prepare outputs
- if return_counts:
- out_data_shape = [data.shape] * 3
- out_dtypes = [data.dtype, "int32", "int32"]
- else:
- out_data_shape = [data.shape] * 2
- out_dtypes = [data.dtype, "int32"]
- # prepare inputs and fcompute
-
- first_occurence = _calc_first_occurence(argsorted_indices, inc_scan)
- if is_sorted:
- in_data = [data, argsorted_indices, inc_scan]
- if return_counts:
- fcompute = lambda ins, outs: _calc_unique_ir(*ins, None, *outs)
- else:
- fcompute = lambda ins, outs: _calc_unique_ir(*ins, None, *outs, None)
-
- indices = first_occurence
- else:
- # calculate index converter by sorting unique elements by their first occurence
- argsorted_first_occurence = argsort(first_occurence, dtype="int32")
- index_converter = argsort(argsorted_first_occurence, dtype="int32")
- in_data = [data, argsorted_indices, inc_scan, index_converter]
- if return_counts:
- fcompute = lambda ins, outs: _calc_unique_ir(*ins, *outs)
- else:
- fcompute = lambda ins, outs: _calc_unique_ir(*ins, *outs, None)
- # First occurence is in order of sorted unique output, if we sort the first_occurence array
- # we get the correct result
- indices = sort(first_occurence)
-
- outs = te.extern(
- out_data_shape,
- in_data,
- fcompute,
- dtype=out_dtypes,
- name="_calc_unique",
- tag="_calc_unique_cpu",
- )
- if return_counts:
- return [outs[0], indices, outs[1], num_unique_elements, outs[2]]
- return [outs[0], indices, outs[1], num_unique_elements]
diff --git a/python/tvm/topi/vision/__init__.py b/python/tvm/topi/vision/__init__.py
deleted file mode 100644
index 2861d31de0..0000000000
--- a/python/tvm/topi/vision/__init__.py
+++ /dev/null
@@ -1,25 +0,0 @@
-# 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.
-
-# pylint: disable=wildcard-import
-"""VISION network operators"""
-from __future__ import absolute_import as _abs
-
-from . import ssd
-from .reorg import *
-from .nms import *
-from .rcnn import *
diff --git a/python/tvm/topi/vision/nms.py b/python/tvm/topi/vision/nms.py
deleted file mode 100644
index 7bd94745e2..0000000000
--- a/python/tvm/topi/vision/nms.py
+++ /dev/null
@@ -1,1183 +0,0 @@
-# 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.
-# pylint: disable=import-error, invalid-name, no-member, too-many-locals, too-many-arguments, undefined-variable, too-many-nested-blocks, too-many-branches, too-many-statements, too-many-function-args
-"""Non-maximum suppression operator"""
-import tvm
-from tvm import te
-
-from tvm.te import hybrid
-from tvm.tir import if_then_else
-
-from ..sort import argsort
-from ..math import cast
-from ..transform import reshape, gather
-from .. import reduction
-from ..scan import cumsum
-from .nms_util import (
- binary_search,
- collect_selected_indices,
- collect_selected_indices_and_scores,
- run_all_class_nms,
-)
-
-
-@hybrid.script
-def hybrid_rearrange_box_out(data, one, batch_size, num_anchors):
- """Hybrid routine to rearrange nms output to
- move all valid entries to top.
-
- Parameters
- ----------
- data : tvm.te.Tensor or numpy NDArray
- NMS output. 3-D tensor with shape
- [batch_size, num_anchors, 6].
-
- one: tvm.tir.const
- Constant one with the same dtype as data.
-
- batch_size: tvm.tir.IntImm or tvm.tir.Var
- Batch size. We need to pass it in since hybrid script doesn't support
- binding variable to symbolic dim.
-
- num_anchors: tvm.tir.IntImm or tvm.tir.Var
- Number of anchors.
-
- Returns
- -------
- output : tvm.te.Tensor or numpy NDArray
- Transformed NMS output. 3-D tensor with shape
- [batch_size, num_anchors, 6].
- """
- elem_length = data.shape[2]
- output = output_tensor((batch_size, num_anchors, elem_length), data.dtype)
- valid_indices = allocate((batch_size,), "int32")
-
- for i in parallel(batch_size):
- valid_indices[i] = 0
- for j in range(num_anchors):
- if data[i, j, 0] >= 0:
- for k in range(elem_length):
- output[i, valid_indices[i], k] = data[i, j, k]
- valid_indices[i] += 1
- if j >= valid_indices[i]:
- for k in range(elem_length):
- output[i, j, k] = -one
- return output
-
-
-@hybrid.script
-def hybrid_rearrange_indices_out(data, one, batch_size, num_anchors):
- """Hybrid routine to rearrange nms output to
- move all valid entries to top.
-
- Parameters
- ----------
- data : tvm.te.Tensor or numpy NDArray
- NMS output. 3-D tensor with shape
- [batch_size, num_anchors, 6] or
- [batch_size, num_anchors, 5], or 2-D
- tensor with shape [batch_size, num_anchors].
-
- one: tvm.tir.const
- Constant one with the same dtype as data.
-
- batch_size: tvm.tir.IntImm or tvm.tir.Var
- Batch size. We need to pass it in since hybrid script doesn't support
- binding variable to symbolic dim.
-
- num_anchors: tvm.tir.IntImm or tvm.tir.Var
- Number of anchors.
-
- Returns
- -------
- output : tvm.te.Tensor or numpy NDArray
- 2-D tensor with shape [batch_size, num_anchors].
-
- valid_box_count : tvm.te.Tensor or numpy NDArray
- Tensor with shape [batch_size, 1], indicates
- the valid number of boxes.
- """
- valid_box_count = output_tensor((batch_size, 1), "int32")
- output = output_tensor((batch_size, num_anchors), data.dtype)
- valid_indices = allocate((batch_size,), "int32")
-
- for i in parallel(batch_size):
- valid_indices[i] = 0
- for j in range(num_anchors):
- if data[i, j] >= 0:
- output[i, valid_indices[i]] = data[i, j]
- valid_indices[i] += 1
- if data[i, j] > num_anchors or data[i, j] < -num_anchors:
- output[i, valid_indices[i]] = 0
- valid_indices[i] += 1
- if j >= valid_indices[i]:
- output[i, j] = -one
- valid_box_count[i, 0] = valid_indices[i]
-
- return output, valid_box_count
-
-
-@hybrid.script
-def hybrid_get_valid_counts(
- data, score_threshold, id_index, score_index, one, batch_size, num_anchors
-):
- """Hybrid routine to get valid count of bounding boxes
- given a score threshold. Also moves valid boxes to the
- top of input data.
-
- Parameters
- ----------
- data : tvm.te.Tensor or numpy NDArray
- Input data. 3-D tensor with shape [batch_size, num_anchors, 6]
- or [batch_size, num_anchors, 5].
-
- score_threshold : tvm.te.Tensor
- Lower limit of score for valid bounding boxes.
-
- id_index : tvm.tir.const
- index of the class categories, -1 to disable.
-
- score_index: tvm.tir.const
- Index of the scores/confidence of boxes.
-
- one: tvm.tir.const
- Constant one with the same dtype as data.
-
- batch_size: tvm.tir.IntImm or tvm.tir.Var
- Batch size. We need to pass it in since hybrid script doesn't support
- binding variable to symbolic dim.
-
- num_anchors: tvm.tir.IntImm or tvm.tir.Var
- Number of anchors.
-
- Returns
- -------
- valid_count : tvm.te.Tensor or numpy NDArray
- 1-D tensor for valid number of boxes.
-
- out_tensor : tvm.te.Tensor or numpy NDArray
- Rearranged data tensor.
-
- out_indices: tvm.te.Tensor or numpy NDArray
- Related index in input data.
- """
- box_data_length = data.shape[2]
- valid_count = output_tensor((batch_size,), "int32")
- out_tensor = output_tensor((batch_size, num_anchors, box_data_length), data.dtype)
- out_indices = output_tensor((batch_size, num_anchors), "int32")
- for i in parallel(batch_size):
- valid_count[i] = 0
- for j in range(num_anchors):
- score = data[i, j, score_index]
- if score > score_threshold and (id_index < 0 or data[i, j, id_index] >= 0):
- for k in range(box_data_length):
- out_tensor[i, valid_count[i], k] = data[i, j, k]
- out_indices[i, valid_count[i]] = j
- valid_count[i] += 1
- if j >= valid_count[i]:
- for k in range(box_data_length):
- out_tensor[i, j, k] = -one
- out_indices[i, j] = -1
- return valid_count, out_tensor, out_indices
-
-
-def get_valid_counts(data, score_threshold=0, id_index=0, score_index=1):
- """Get valid count of bounding boxes given a score threshold.
- Also moves valid boxes to the top of input data.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- Input data. 3-D tensor with shape [batch_size, num_anchors, 6]
- or [batch_size, num_anchors, 5].
-
- score_threshold : optional, float
- Lower limit of score for valid bounding boxes.
-
- id_index : optional, int
- index of the class categories, -1 to disable.
-
- score_index: optional, int
- Index of the scores/confidence of boxes.
-
- Returns
- -------
- valid_count : tvm.te.Tensor
- 1-D tensor for valid number of boxes.
-
- out_tensor : tvm.te.Tensor
- Rearranged data tensor.
-
- out_indices: tvm.te.Tensor or numpy NDArray
- Related index in input data.
- """
- if isinstance(score_threshold, (float, int)):
- score_threshold = tvm.tir.const(score_threshold, dtype=data.dtype)
- id_index_const = tvm.tir.const(id_index, "int32")
- score_index_const = tvm.tir.const(score_index, "int32")
- return hybrid_get_valid_counts(
- data,
- score_threshold,
- id_index_const,
- score_index_const,
- tvm.tir.const(1, data.dtype),
- data.shape[0],
- data.shape[1],
- )
-
-
-@hybrid.script
-def hybrid_nms(
- data,
- sorted_index,
- valid_count,
- indices,
- batch_size,
- num_anchors,
- max_output_size,
- iou_threshold,
- force_suppress,
- top_k,
- coord_start,
- score_index,
- id_index,
- return_indices,
- zero,
- one,
-):
- """Hybrid routing for non-maximum suppression.
-
- Parameters
- ----------
- data: tvm.te.Tensor or numpy NDArray
- Bounding boxes with class and score. 3-D tensor with shape
- [batch_size, num_anchors, 6]. It could be the second output
- out_tensor of get_valid_counts.
-
- sorted_index : tvm.te.Tensor or numpy NDArray
- Bounding box indexes sorted by score, with shape
- [batch_size, num_anchors].
-
- valid_count : tvm.te.Tensor or numpy NDArray
- 1-D tensor for valid number of boxes. It could be the output
- valid_count of get_valid_counts.
-
- indices : tvm.te.Tensor or numpy.NDArray
- indices in original tensor, with shape [batch_size, num_anchors],
- represents the index of box in original data. It could be the third
- output out_indices of get_valid_counts. The values in the second
- dimension are like the output of arange(num_anchors) if get_valid_counts
- is not used before non_max_suppression.
-
- batch_size: tvm.tir.IntImm or tvm.tir.Var
- Batch size. We need to pass it in since hybrid script doesn't support
- binding variable to symbolic dim.
-
- num_anchors: tvm.tir.IntImm or tvm.tir.Var
- The number of anchors.
-
- max_output_size : tvm.te.Tensor
- Max number of output valid boxes for each instance.
- Return all valid boxes if max_output_size < 0.
-
- iou_threshold : tvm.te.Tensor
- Overlapping(IoU) threshold to suppress object with smaller score.
-
- force_suppress : tvm.tir.const
- Whether to suppress all detections regardless of class_id.
-
- top_k : tvm.tir.const
- Keep maximum top k detections before nms, -1 for no limit.
-
- coord_start : tvm.tir.const
- Start index of the consecutive 4 coordinates.
-
- score_index: tvm.tir.const
- Index of the scores/confidence of boxes.
-
- id_index : tvm.tir.const
- index of the class categories, -1 to disable.
-
- return_indices : tvm.tir.const
- Whether to return box indices in input data.
-
- zero: tvm.tir.const
- Constant zero with the same dtype as data.
-
- one: tvm.tir.const
- Constant one with the same dtype as data.
-
- Returns
- -------
- output : tvm.te.Tensor
- 3-D tensor with shape [batch_size, num_anchors, 6]
- or [batch_size, num_anchors, 5].
-
- box_indices: tvm.te.Tensor
- 2-D tensor with shape [batch_size, num_anchors].
- """
-
- box_data_length = data.shape[2]
-
- # box_indices is the expected indices of boxes
- box_indices = output_tensor((batch_size, num_anchors), sorted_index.dtype)
- output = output_tensor(
- (
- batch_size,
- num_anchors,
- box_data_length,
- ),
- data.dtype,
- )
-
- for i in range(batch_size):
- if iou_threshold > 0:
- if valid_count[i] > 0:
- # Reorder output
- nkeep = valid_count[i]
- if 0 < top_k < nkeep:
- nkeep = top_k
- for j in parallel(nkeep):
- for k in range(box_data_length):
- output[i, j, k] = data[i, sorted_index[i, j], k]
- box_indices[i, j] = sorted_index[i, j]
- if 0 < top_k < valid_count[i]:
- for j in parallel(valid_count[i] - nkeep):
- for k in range(box_data_length):
- output[i, j + nkeep, k] = -one
- box_indices[i, j + nkeep] = -1
-
- # Apply nms
- box_start_idx = coord_start
- batch_idx = i
- num_valid_boxes = 0
-
- for j in range(valid_count[i]):
- if num_valid_boxes == max_output_size:
- for k in range(box_data_length):
- output[i, j, k] = -one
- box_indices[i, j] = -1
-
- elif output[i, j, score_index] > 0:
- box_a_idx = j
- is_valid_box = 1
-
- # a_l: left, a_t: top, a_r: right, a_b: bottom
- a_l = min(
- output[batch_idx, box_a_idx, box_start_idx],
- output[batch_idx, box_a_idx, box_start_idx + 2],
- )
- a_t = min(
- output[batch_idx, box_a_idx, box_start_idx + 1],
- output[batch_idx, box_a_idx, box_start_idx + 3],
- )
- a_r = max(
- output[batch_idx, box_a_idx, box_start_idx],
- output[batch_idx, box_a_idx, box_start_idx + 2],
- )
- a_b = max(
- output[batch_idx, box_a_idx, box_start_idx + 1],
- output[batch_idx, box_a_idx, box_start_idx + 3],
- )
-
- # check if current box j is valid by calculating iou with
- # all existing valid boxes
- for k in range(j):
- check_iou = 0
- if (
- is_valid_box == 1
- and k < j
- and output[i, k, score_index] > 0
- and (id_index < 0 or output[i, k, id_index] >= 0)
- ):
- if force_suppress:
- check_iou = 1
- elif id_index < 0 or output[i, j, id_index] == output[i, k, id_index]:
- check_iou = 1
-
- if check_iou > 0:
- box_b_idx = k
-
- # b_l: left, b_t: top, b_r: right, b_b: bottom
- b_l = min(
- output[batch_idx, box_b_idx, box_start_idx],
- output[batch_idx, box_b_idx, box_start_idx + 2],
- )
- b_t = min(
- output[batch_idx, box_b_idx, box_start_idx + 1],
- output[batch_idx, box_b_idx, box_start_idx + 3],
- )
- b_r = max(
- output[batch_idx, box_b_idx, box_start_idx],
- output[batch_idx, box_b_idx, box_start_idx + 2],
- )
- b_b = max(
- output[batch_idx, box_b_idx, box_start_idx + 1],
- output[batch_idx, box_b_idx, box_start_idx + 3],
- )
-
- # Overlapping width and height
- w = max(zero, min(a_r, b_r) - max(a_l, b_l))
- h = max(zero, min(a_b, b_b) - max(a_t, b_t))
-
- # Overlapping area
- area = h * w
-
- # total area of the figure formed by box a and box b
- # except for overlapping area
- u = (a_r - a_l) * (a_b - a_t) + (b_r - b_l) * (b_b - b_t) - area
-
- # get the iou
- iou = zero if u <= zero else area / u
-
- if iou >= iou_threshold:
- is_valid_box = 0
-
- if is_valid_box == 0:
- for k in range(box_data_length):
- output[i, j, k] = -one
- box_indices[i, j] = -1
- else:
- num_valid_boxes += 1
-
- else:
- for j in parallel(valid_count[i]):
- for k in range(box_data_length):
- output[i, j, k] = data[i, j, k]
- box_indices[i, j] = j
-
- # Set invalid entry to be -1
- for j in parallel(num_anchors - valid_count[i]):
- for k in range(box_data_length):
- output[i, j + valid_count[i], k] = -one
- box_indices[i, j + valid_count[i]] = -1
-
- if return_indices:
- for j in range(valid_count[i]):
- idx = box_indices[i, j]
- if box_indices[i, j] >= 0:
- box_indices[i, j] = indices[i, idx]
-
- return output, box_indices
-
-
-@tvm.target.generic_func
-def non_max_suppression(
- data,
- valid_count,
- indices,
- max_output_size=-1,
- iou_threshold=0.5,
- force_suppress=False,
- top_k=-1,
- coord_start=2,
- score_index=1,
- id_index=0,
- return_indices=True,
- invalid_to_bottom=False,
-):
- """Non-maximum suppression operator for object detection.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- 3-D tensor with shape [batch_size, num_anchors, 6] or [batch_size, num_anchors, 5].
-
- valid_count : tvm.te.Tensor
- 1-D tensor for valid number of boxes.
-
- indices : tvm.te.Tensor
- 2-D tensor with shape [batch_size, num_anchors].
-
- max_output_size : optional, int or tvm.te.Tensor
- Max number of output valid boxes for each instance.
- Return all valid boxes if the value of max_output_size is less than 0.
-
- iou_threshold : optional, float or tvm.te.Tensor
- Non-maximum suppression threshold.
-
- force_suppress : optional, boolean
- Whether to suppress all detections regardless of class_id.
-
- top_k : optional, int
- Keep maximum top k detections before nms, -1 for no limit.
-
- coord_start : required, int
- Start index of the consecutive 4 coordinates.
-
- score_index: optional, int
- Index of the scores/confidence of boxes.
-
- id_index : optional, int
- index of the class categories, -1 to disable.
-
- return_indices : optional, boolean
- Whether to return box indices in input data.
-
- invalid_to_bottom : optional, boolean
- Whether to move all valid bounding boxes to the top.
-
- Returns
- -------
- out : tvm.te.Tensor or tuple of tvm.te.Tensor
- 3-D tensor with shape [batch_size, num_anchors, 6]
- or [batch_size, num_anchors, 5]. Out is a tuple of tvm.te.Tensor
- if return_indices is True, the Tensor in the tuple is 2-D tensor
- with shape [batch_size, num_anchors] and shape
- [batch_size, num_valid_anchors] respectively.
-
- Example
- --------
- .. code-block:: python
-
- # An example to use non_max_suppression
- dshape = (1, 5, 6)
- data = te.placeholder(dshape, name="data")
- valid_count = te.placeholder((dshape[0],), dtype="int32", name="valid_count")
- iou_threshold = 0.7
- force_suppress = True
- top_k = -1
- out = non_max_suppression(data, valid_count, indices, iou_threshold=iou_threshold,
- force_suppress=force_suppress, top_k=top_k)
- np_data = np.random.uniform(dshape)
- np_valid_count = np.array([4])
- s = topi.generic.schedule_nms(out)
- f = tvm.build(s, [data, valid_count, out], "llvm")
- dev = tvm.cpu()
- tvm_data = tvm.nd.array(np_data, dev)
- tvm_valid_count = tvm.nd.array(np_valid_count, dev)
- tvm_out = tvm.nd.array(np.zeros(dshape, dtype=data.dtype), dev)
- f(tvm_data, tvm_valid_count, tvm_out)
- """
- batch_size = data.shape[0]
- num_anchors = data.shape[1]
- if isinstance(max_output_size, int):
- max_output_size = tvm.tir.const(max_output_size, dtype="int32")
- if isinstance(iou_threshold, float):
- iou_threshold = tvm.tir.const(iou_threshold, dtype=data.dtype)
- score_axis = score_index
- score_shape = (batch_size, num_anchors)
- score_tensor = te.compute(score_shape, lambda i, j: data[i, j, score_axis])
- sort_tensor = argsort(score_tensor, valid_count=valid_count, axis=1, is_ascend=False)
-
- out, box_indices = hybrid_nms(
- data,
- sort_tensor,
- valid_count,
- indices,
- batch_size,
- num_anchors,
- max_output_size,
- iou_threshold,
- tvm.tir.const(force_suppress, dtype="bool"),
- tvm.tir.const(top_k, dtype="int32"),
- tvm.tir.const(coord_start, dtype="int32"),
- tvm.tir.const(score_index, dtype="int32"),
- tvm.tir.const(id_index, dtype="int32"),
- tvm.tir.const(return_indices, dtype="bool"),
- zero=tvm.tir.const(0, dtype=data.dtype),
- one=tvm.tir.const(1, dtype=data.dtype),
- )
-
- if return_indices:
- return hybrid_rearrange_indices_out(
- box_indices,
- one=tvm.tir.const(1, dtype="int32"),
- batch_size=batch_size,
- num_anchors=num_anchors,
- )
-
- if invalid_to_bottom:
- out = hybrid_rearrange_box_out(
- out,
- one=tvm.tir.const(1, dtype=data.dtype),
- batch_size=batch_size,
- num_anchors=num_anchors,
- )
- return out
-
-
-def _nms_loop(
- ib,
- batch_size,
- top_k,
- iou_threshold,
- max_output_size,
- valid_count,
- on_new_valid_box_func,
- on_new_invalidated_box_func,
- needs_bbox_check_func,
- calc_overlap_func,
- out_scores,
- num_valid_boxes,
-):
- def nms_inner_loop(ib, i, j, nkeep, num_valid_boxes_local):
- # The box j is valid, invalidate other boxes that overlap with j above iou_threshold
- on_new_valid_box_func(ib, 0, num_valid_boxes_local[0], i, j)
- num_valid_boxes_local[0] += 1
-
- num_boxes_to_check = nkeep - (j + 1)
-
- with ib.for_range(0, num_boxes_to_check, name="_k", kind="parallel") as _k:
- k = j + 1 + _k
-
- with ib.if_scope(
- tvm.tir.all(
- k < nkeep,
- out_scores[i, k] > 0, # is the box k still valid?
- needs_bbox_check_func(i, j, k),
- )
- ):
- iou = calc_overlap_func(i, j, k)
-
- with ib.if_scope(iou >= iou_threshold):
- # invalidate the box k
- out_scores[i, k] = -1.0
- on_new_invalidated_box_func(i, k)
-
- with ib.for_range(0, batch_size, name="i") as i:
- nkeep = if_then_else(tvm.tir.all(top_k > 0, top_k < valid_count[i]), top_k, valid_count[i])
- max_output_size = if_then_else(max_output_size > 0, max_output_size, nkeep)
-
- with ib.if_scope(tvm.tir.all(iou_threshold > 0, valid_count[i] > 0)):
- num_valid_boxes_local = ib.allocate(
- "int32", (1,), name="num_valid_boxes_local", scope="local"
- )
- box_idx = ib.allocate("int32", (1,), name="box_idx", scope="local")
- num_valid_boxes_local[0] = 0
- box_idx[0] = 0
-
- # Apply nms
- # No need to do more iteration if we have already reached max_output_size boxes
- with ib.while_loop(
- tvm.tir.all(box_idx[0] < nkeep, num_valid_boxes_local[0] < max_output_size)
- ):
- # Proceed to the inner loop if the box with id box_idx is still valid
- with ib.if_scope(out_scores[i, box_idx[0]] > -1.0):
- nms_inner_loop(ib, i, box_idx[0], nkeep, num_valid_boxes_local)
- box_idx[0] += 1
-
- num_valid_boxes[i] = num_valid_boxes_local[0]
-
- with ib.else_scope():
- num_valid_boxes[i] = 0
-
- return ib.get()
-
-
-def _get_valid_box_count(scores, score_threshold):
- batch_classes, num_boxes = scores.shape
-
- def searchsorted_ir(scores, valid_count):
- ib = tvm.tir.ir_builder.create()
- scores = ib.buffer_ptr(scores)
- valid_count = ib.buffer_ptr(valid_count)
-
- with ib.for_range(0, batch_classes, name="i", kind="parallel") as i:
- binary_search(ib, i, num_boxes, scores, score_threshold, valid_count)
-
- return ib.get()
-
- scores_buf = tvm.tir.decl_buffer(scores.shape, scores.dtype, "scores_buf", data_alignment=8)
-
- return te.extern(
- [(batch_classes,)],
- [scores],
- lambda ins, outs: searchsorted_ir(ins[0], outs[0]),
- dtype=["int32"],
- in_buffers=[scores_buf],
- name="searchsorted",
- tag="searchsorted",
- )
-
-
-def _collect_selected_indices_ir(num_class, selected_indices, num_detections, row_offsets, out):
- batch_classes, _ = selected_indices.shape
-
- ib = tvm.tir.ir_builder.create()
-
- selected_indices = ib.buffer_ptr(selected_indices)
- num_detections = ib.buffer_ptr(num_detections)
- row_offsets = ib.buffer_ptr(row_offsets)
- out = ib.buffer_ptr(out)
-
- with ib.for_range(0, batch_classes, name="i", kind="parallel") as i:
- i = cast(i, "int64")
- batch_id = i // num_class
- class_id = i % num_class
-
- with ib.for_range(0, num_detections[i], name="j") as j:
- out[row_offsets[i] + j, 0] = batch_id
- out[row_offsets[i] + j, 1] = class_id
- out[row_offsets[i] + j, 2] = cast(selected_indices[i, j], "int64")
-
- return ib.get()
-
-
-def _collect_selected_indices_and_scores_ir(
- selected_indices,
- selected_scores,
- num_detections,
- row_offsets,
- num_total_detections,
- collected_indices,
- collected_scores,
-):
- batch_size, num_class = row_offsets.shape
- num_boxes = selected_indices.shape[1]
-
- ib = tvm.tir.ir_builder.create()
-
- selected_indices = ib.buffer_ptr(selected_indices)
- selected_scores = ib.buffer_ptr(selected_scores)
- num_detections = ib.buffer_ptr(num_detections)
- row_offsets = ib.buffer_ptr(row_offsets)
- num_total_detections = ib.buffer_ptr(num_total_detections)
- collected_indices = ib.buffer_ptr(collected_indices)
- collected_scores = ib.buffer_ptr(collected_scores)
- zero = cast(0, "int64")
-
- with ib.for_range(0, batch_size * num_class, name="i", kind="parallel") as i:
- i = cast(i, "int64")
- batch_id = i // num_class
- class_id = i % num_class
-
- with ib.for_range(0, num_boxes, name="j") as j:
- with ib.if_scope(j < num_detections[batch_id, class_id]):
- offset = row_offsets[batch_id, class_id] + j
- collected_indices[batch_id, offset, 0] = class_id
- collected_indices[batch_id, offset, 1] = cast(selected_indices[i, j], "int64")
- collected_scores[batch_id, offset] = selected_scores[i, j]
- with ib.else_scope():
- offset = (
- num_total_detections[batch_id]
- + class_id * num_boxes
- - row_offsets[batch_id, class_id]
- + j
- - num_detections[batch_id, class_id]
- )
- collected_indices[batch_id, offset, 0] = zero
- collected_indices[batch_id, offset, 1] = zero
- collected_scores[batch_id, offset] = 0.0
-
- return ib.get()
-
-
-def all_class_non_max_suppression(
- boxes,
- scores,
- max_output_boxes_per_class,
- iou_threshold,
- score_threshold,
- output_format="onnx",
-):
- """Non-maximum suppression operator for object detection, corresponding to ONNX
- NonMaxSuppression and TensorFlow combined_non_max_suppression.
- NMS is performed for each class separately.
-
- Parameters
- ----------
- boxes : tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_boxes, 4)
-
- scores: tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_classes, num_boxes)
-
- max_output_boxes_per_class : int or tvm.te.Tensor, optional
- The maxinum number of output selected boxes per class
-
- iou_threshold : float or tvm.te.Tensor, optionaIl
- IoU test threshold
-
- score_threshold : float or tvm.te.Tensor, optional
- Score threshold to filter out low score boxes early
-
- output_format : str, optional
- "onnx" or "tensorflow", see below.
-
- Returns
- -------
- out : list of tvm.te.Tensor
- If `output_format` is "onnx", the output is two tensors. The first is `indices` of size
- `(batch_size * num_class* num_boxes , 3)` and the second is a scalar tensor
- `num_total_detection` of shape `(1,)` representing the total number of selected
- boxes. The three values in `indices` encode batch, class, and box indices.
- Rows of `indices` are ordered such that selected boxes from batch 0, class 0 come
- first, in descending of scores, followed by boxes from batch 0, class 1 etc. Out of
- `batch_size * num_class* num_boxes` rows of indices, only the first `num_total_detection`
- rows are valid.
-
- If `output_format` is "tensorflow", the output is three tensors, the first
- is `indices` of size `(batch_size, num_class * num_boxes , 2)`, the second is `scores` of
- size `(batch_size, num_class * num_boxes)`, and the third is `num_total_detection` of size
- `(batch_size,)` representing the total number of selected boxes per batch. The two values
- in `indices` encode class and box indices. Of num_class * num_boxes boxes in `indices` at
- batch b, only the first `num_total_detection[b]` entries are valid. The second axis of
- `indices` and `scores` are sorted within each class by box scores, but not across classes.
- So the box indices and scores for the class 0 come first in a sorted order, followed by
- the class 1 etc.
- """
- batch, num_class, num_boxes = scores.shape
- scores = reshape(scores, (batch * num_class, num_boxes))
-
- sorted_indices = argsort(scores, axis=1, is_ascend=False, dtype="int32")
- sorted_scores = gather(scores, 1, sorted_indices)
-
- valid_count = _get_valid_box_count(sorted_scores, score_threshold)
-
- selected_indices, selected_scores, num_detections = run_all_class_nms(
- boxes,
- sorted_scores,
- sorted_indices,
- valid_count,
- max_output_boxes_per_class,
- iou_threshold,
- _nms_loop,
- return_scores=(output_format == "tensorflow"),
- )
-
- if output_format == "onnx":
- row_offsets = cumsum(num_detections, exclusive=True, dtype="int64")
- num_total_detections = reduction.sum(cast(num_detections, "int64"), axis=1)
-
- selected_indices = collect_selected_indices(
- num_class, selected_indices, num_detections, row_offsets, _collect_selected_indices_ir
- )
- return [selected_indices, num_total_detections]
-
- num_detections_per_batch = reshape(num_detections, (batch, num_class))
- row_offsets = cumsum(num_detections_per_batch, exclusive=True, dtype="int64", axis=1)
- num_total_detections = reduction.sum(cast(num_detections_per_batch, "int64"), axis=1)
-
- selected_indices, selected_scores = collect_selected_indices_and_scores(
- selected_indices,
- selected_scores,
- num_detections_per_batch,
- row_offsets,
- num_total_detections,
- _collect_selected_indices_and_scores_ir,
- )
-
- return [selected_indices, selected_scores, num_total_detections]
-
-
-@hybrid.script
-def hybrid_regular_nms(
- boxes,
- scores,
- max_detections_per_class,
- max_detections,
- batch_size,
- num_boxes,
- num_classes,
- num_classes_with_background,
- iou_threshold,
- score_threshold,
-):
- """Hybrid routing for regular non-maximum suppression.
-
- Parameters
- ----------
- boxes : tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_boxes, 4)
-
- scores: tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_boxes, num_classes_with_background)
-
- max_detections_per_class : tvm.tir.const
- The maxinum number of output selected boxes per class
-
- max_detections : tvm.tir.const
- The maxinum number of output selected boxes
-
- batch_size : tvm.tir.IntImm or tvm.tir.Var
- The number of batches
-
- num_boxes : tvm.tir.IntImm or tvm.tir.Var
- The number of bounding boxes
-
- num_classes : tvm.tir.const
- The number of classes without background
-
- num_classes_with_background : tvm.tir.IntImm or tvm.tir.Var
- The number of classes including background ones
-
- iou_threshold : tvm.tir.const
- IoU test threshold
-
- score_threshold : tvm.tir.const
- Score threshold to filter out low score boxes early
-
- Returns
- -------
- detection_boxes : tvm.te.Tensor
- 3-D tensor with shape [batch_size, max_detections, 4].
-
- detection_classes : tvm.te.Tensor
- 2-D tensor with shape [batch_size, max_detections].
-
- detection_scores : tvm.te.Tensor
- 2-D tensor with shape [batch_size, max_detections].
-
- num_detections : tvm.te.Tensor
- 1-D tensor with shape [batch_size].
- """
- # output tensors
- detection_boxes = output_tensor((batch_size, max_detections, 4), boxes.dtype)
- detection_classes = output_tensor((batch_size, max_detections), "int32")
- detection_scores = output_tensor((batch_size, max_detections), scores.dtype)
- num_detections = output_tensor((batch_size,), "int32")
-
- # scratch buffers
- class_scores = allocate((num_boxes,), scores.dtype)
- keep_indices = allocate((num_boxes,), "int32")
- keep_scores = allocate((num_boxes,), scores.dtype)
- sorted_indices = allocate((max_detections + num_boxes,), "int32")
- sorted_scores = allocate((max_detections + num_boxes,), scores.dtype)
- active_box_candidate = allocate((num_boxes,), "int32")
- selected = allocate((num_boxes,), "int32")
- box_indices_after_regular_nms = allocate((max_detections + num_boxes,), "int32")
- scores_after_regular_nms = allocate((max_detections + num_boxes,), scores.dtype)
-
- label_offset = num_classes_with_background - num_classes
- tmp_idx = 0
-
- for batch_idx in range(batch_size):
- size_of_sorted_indices = 0
-
- for class_id in range(num_classes):
- for box_id in range(num_boxes):
- # get scores of boxes corresponding to all anchors for single class
- class_scores[box_id] = scores[batch_idx, box_id, class_id + label_offset]
-
- # perform non-maximal suppression on single class
-
- # select detections above score threshold
- num_scores_kept = 0
- for i in range(num_boxes):
- if class_scores[i] >= score_threshold:
- keep_scores[num_scores_kept] = class_scores[i]
- keep_indices[num_scores_kept] = i
- num_scores_kept += 1
-
- # iota
- for i in range(num_scores_kept):
- sorted_indices[i] = i
- # decreasing sort of scores
- for i in range(num_scores_kept):
- for j in range(num_scores_kept - i - 1):
- if keep_scores[sorted_indices[j]] < keep_scores[sorted_indices[j + 1]]:
- tmp_idx = sorted_indices[j]
- sorted_indices[j] = sorted_indices[j + 1]
- sorted_indices[j + 1] = tmp_idx
-
- selected_size = 0
-
- for i in range(num_scores_kept):
- active_box_candidate[i] = 1
-
- num_active_candidate = num_scores_kept
- for i in range(num_scores_kept):
- if (
- num_active_candidate != 0
- and selected_size < min(num_scores_kept, max_detections_per_class)
- and active_box_candidate[i] == 1
- ):
- selected[selected_size] = keep_indices[sorted_indices[i]]
- selected_size += 1
-
- active_box_candidate[i] = 0
- num_active_candidate -= 1
-
- for j in range(i + 1, num_scores_kept):
- if active_box_candidate[j] == 1:
- # compute IOU
- i_ymin = boxes[batch_idx, keep_indices[sorted_indices[i]], 0]
- i_xmin = boxes[batch_idx, keep_indices[sorted_indices[i]], 1]
- i_ymax = boxes[batch_idx, keep_indices[sorted_indices[i]], 2]
- i_xmax = boxes[batch_idx, keep_indices[sorted_indices[i]], 3]
-
- j_ymin = boxes[batch_idx, keep_indices[sorted_indices[j]], 0]
- j_xmin = boxes[batch_idx, keep_indices[sorted_indices[j]], 1]
- j_ymax = boxes[batch_idx, keep_indices[sorted_indices[j]], 2]
- j_xmax = boxes[batch_idx, keep_indices[sorted_indices[j]], 3]
-
- area_i = (i_ymax - i_ymin) * (i_xmax - i_xmin)
- area_j = (j_ymax - j_ymin) * (j_xmax - j_xmin)
-
- iou = 0.0
- if area_i > 0 and area_j > 0:
- intersection_ymin = max(i_ymin, j_ymin)
- intersection_xmin = max(i_xmin, j_xmin)
- intersection_ymax = min(i_ymax, j_ymax)
- intersection_xmax = min(i_xmax, j_xmax)
- intersection_area = max(
- intersection_ymax - intersection_ymin, 0.0
- ) * max(intersection_xmax - intersection_xmin, 0.0)
- iou = intersection_area / (area_i + area_j - intersection_area)
-
- if iou > iou_threshold:
- active_box_candidate[j] = 0
- num_active_candidate -= 1
-
- # end of non-maximal suppression on single class
-
- # add selected indices from non-max suppression of boxes in this class
- output_index = size_of_sorted_indices
- for i in range(selected_size):
- selected_index = selected[i]
-
- box_indices_after_regular_nms[output_index] = (
- selected_index * num_classes_with_background + class_id + label_offset
- )
- scores_after_regular_nms[output_index] = class_scores[selected_index]
-
- output_index += 1
-
- # sort the max scores among the selected indices
- # get the indices for top scores
- num_indices_to_sort = min(output_index, max_detections)
-
- # iota
- for i in range(output_index):
- sorted_indices[i] = i
- # deacreasing sort of scores
- for i in range(output_index):
- for j in range(output_index - i - 1):
- if (
- scores_after_regular_nms[sorted_indices[j]]
- < scores_after_regular_nms[sorted_indices[j + 1]]
- ):
- tmp_idx = sorted_indices[j]
- sorted_indices[j] = sorted_indices[j + 1]
- sorted_indices[j + 1] = tmp_idx
-
- # copy values to temporary vectors
- for i in range(num_indices_to_sort):
- sorted_scores[i] = scores_after_regular_nms[sorted_indices[i]]
- sorted_indices[i] = box_indices_after_regular_nms[sorted_indices[i]]
-
- # copy scores and indices from temporary vectors
- for i in range(num_indices_to_sort):
- box_indices_after_regular_nms[i] = sorted_indices[i]
- scores_after_regular_nms[i] = sorted_scores[i]
-
- size_of_sorted_indices = num_indices_to_sort
-
- # fill output tensors
- for output_box_index in range(max_detections):
- box_ymin = 0.0
- box_xmin = 0.0
- box_ymax = 0.0
- box_xmax = 0.0
- class_idx = 0
- selected_score = 0.0
-
- if output_box_index < size_of_sorted_indices:
- anchor_idx = (
- box_indices_after_regular_nms[output_box_index] // num_classes_with_background
- )
-
- box_ymin = boxes[batch_idx, anchor_idx, 0]
- box_xmin = boxes[batch_idx, anchor_idx, 1]
- box_ymax = boxes[batch_idx, anchor_idx, 2]
- box_xmax = boxes[batch_idx, anchor_idx, 3]
- class_idx = (
- box_indices_after_regular_nms[output_box_index]
- - anchor_idx * num_classes_with_background
- - label_offset
- )
- selected_score = scores_after_regular_nms[output_box_index]
-
- detection_boxes[batch_idx, output_box_index, 0] = box_ymin
- detection_boxes[batch_idx, output_box_index, 1] = box_xmin
- detection_boxes[batch_idx, output_box_index, 2] = box_ymax
- detection_boxes[batch_idx, output_box_index, 3] = box_xmax
- detection_classes[batch_idx, output_box_index] = class_idx
- detection_scores[batch_idx, output_box_index] = selected_score
-
- num_detections[batch_idx] = size_of_sorted_indices
-
- return detection_boxes, detection_classes, detection_scores, num_detections
-
-
-def regular_non_max_suppression(
- boxes,
- scores,
- max_detections_per_class,
- max_detections,
- num_classes,
- iou_threshold,
- score_threshold,
-):
- """Regular non-maximum suppression operator for object detection, corresponding to TFLite's
- regular NMS. NMS is performed for each class separately.
-
- Parameters
- ----------
- boxes : tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_boxes, 4). The four values in boxes
- encode (ymin, xmin, ymax, xmax) coordinates of a box
-
- scores: tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_boxes, num_classes_with_background)
-
- max_detections_per_class : int
- The maxinum number of output selected boxes per class
-
- max_detections : int
- The maxinum number of output selected boxes
-
- num_classes : int
- The number of classes without background
-
- iou_threshold : float
- IoU test threshold
-
- score_threshold : float
- Score threshold to filter out low score boxes early
-
- Returns
- -------
- out : list of tvm.te.Tensor
- The output is a list of four tensors. The first is `detection_boxes` of size
- `(batch_size, max_detections , 4)`, the second is `detection_classes` of size
- `(batch_size, max_detections)`, the third is `detection_scores` of size
- `(batch_size, max_detections)`, and the fourth is `num_detections` of size `(batch_size,)`
- representing the total number of selected boxes per batch.
- """
- batch_size, num_boxes, num_classes_with_background = scores.shape
-
- detection_boxes, detection_classes, detection_scores, num_detections = hybrid_regular_nms(
- boxes=boxes,
- scores=scores,
- max_detections_per_class=tvm.tir.const(max_detections_per_class, dtype="int32"),
- max_detections=tvm.tir.const(max_detections, dtype="int32"),
- batch_size=batch_size,
- num_boxes=num_boxes,
- num_classes=tvm.tir.const(num_classes, dtype="int32"),
- num_classes_with_background=num_classes_with_background,
- iou_threshold=tvm.tir.const(iou_threshold, dtype="float32"),
- score_threshold=tvm.tir.const(score_threshold, dtype="float32"),
- )
-
- return [
- detection_boxes,
- cast(detection_classes, dtype="float32"),
- detection_scores,
- num_detections,
- ]
diff --git a/python/tvm/topi/vision/nms_util.py b/python/tvm/topi/vision/nms_util.py
deleted file mode 100644
index d12592fd11..0000000000
--- a/python/tvm/topi/vision/nms_util.py
+++ /dev/null
@@ -1,338 +0,0 @@
-# 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.
-# pylint: disable=invalid-name
-"""Common utilities used in Non-maximum suppression operators"""
-import tvm
-from tvm import te
-
-
-def _get_boundaries(output, box_idx):
- l = tvm.te.min(
- output[box_idx],
- output[box_idx + 2],
- )
- t = tvm.te.min(
- output[box_idx + 1],
- output[box_idx + 3],
- )
- r = tvm.te.max(
- output[box_idx],
- output[box_idx + 2],
- )
- b = tvm.te.max(
- output[box_idx + 1],
- output[box_idx + 3],
- )
- return l, t, r, b
-
-
-def calculate_overlap(out_tensor, box_a_idx, box_b_idx):
- """Calculate overlap of two boxes."""
- a_l, a_t, a_r, a_b = _get_boundaries(out_tensor, box_a_idx)
- b_l, b_t, b_r, b_b = _get_boundaries(out_tensor, box_b_idx)
-
- # Overlapping width and height
- w = tvm.te.max(0.0, tvm.te.min(a_r, b_r) - tvm.te.max(a_l, b_l))
- h = tvm.te.max(0.0, tvm.te.min(a_b, b_b) - tvm.te.max(a_t, b_t))
-
- # Overlapping area
- area = h * w
-
- # total area of the figure formed by box a and box b
- # except for overlapping area
- u = (a_r - a_l) * (a_b - a_t) + (b_r - b_l) * (b_b - b_t) - area
- return tvm.tir.Select(u <= 0.0, 0.0, area / u)
-
-
-def binary_search(ib, y, num_boxes, scores, score_threshold, out):
- """Binary search for score_threshold on scores sorted in descending order"""
- lo = ib.allocate("int32", (1,), name="lo", scope="local")
- hi = ib.allocate("int32", (1,), name="hi", scope="local")
-
- lo[0] = 0
- hi[0] = num_boxes
-
- with ib.while_loop(lo[0] < hi[0]):
- mid = (hi[0] + lo[0]) >> 1
- with ib.if_scope(scores[y, mid] > score_threshold):
- lo[0] = mid + 1
- with ib.else_scope():
- hi[0] = mid
-
- out[y] = lo[0]
-
-
-def collect_selected_indices(num_class, selected_indices, num_detections, row_offsets, ir):
- """Collect selected indices from the core NMS loop into one linear output
-
- Parameters
- ----------
- num_class : int
-
- selected_indices: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
- of selected boxes by the core NMS loop.
-
- num_detections tvm.te.Tensor
- 1-D tensor with shape (batch_size * num_classes,), representing
- the number of boxes selected by the core NMS loop, per batch and class
-
- row_offsets tvm.te.Tensor
- 1-D tensor with shape (batch_size * num_classes,), this should be the exclusive scan
- of num_detections
-
- ir : function
- A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py
-
- Returns
- -------
- out : tvm.te.Tensor
- The output is indices of size (batch_size * num_class* num_boxes , 3).
- Rows of indices are ordered such that selected boxes from batch 0, class 0 come
- first, in descending of scores, followed by boxes from batch 0, class 1 etc.
- """
- batch_class, num_boxes = selected_indices.shape
- return te.extern(
- [(batch_class * num_boxes, 3)],
- [selected_indices, num_detections, row_offsets],
- lambda ins, outs: ir(num_class, ins[0], ins[1], ins[2], outs[0]),
- dtype=["int64"],
- name="collect_indices",
- tag="collect_indices",
- )
-
-
-def collect_selected_indices_and_scores(
- selected_indices, selected_scores, num_detections, row_offsets, num_total_detections, ir
-):
- """Collect selected indices and scores from the core NMS loop into one linear output
-
- Parameters
- ----------
- num_class : int
-
- selected_indices: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
- of selected boxes by the core NMS loop.
-
- selected_indices: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes), representing the scores
- of selected boxes by the core NMS loop.
-
- num_detections tvm.te.Tensor
- 2-D tensor with shape (batch_size, num_classes), representing
- the number of boxes selected by the core NMS loop, per batch and class
-
- row_offsets tvm.te.Tensor
- 2-D tensor with shape (batch_size, num_classes), this should be the exclusive scan
- of num_detections along axis 1
-
- ir : function
- A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py
-
- Returns
- -------
- out : [tvm.te.Tensor, tvm.te.Tensor]
- The output is two tensors. The first is indices of size
- (batch_size, num_class* num_boxes, 2), and the second is scores of size
- (batch_size, num_class* num_boxes).
- """
- batch_size, num_class = row_offsets.shape
- num_boxes = selected_indices.shape[1]
- return te.extern(
- [(batch_size, num_class * num_boxes, 2), (batch_size, num_class * num_boxes)],
- [selected_indices, selected_scores, num_detections, row_offsets, num_total_detections],
- lambda ins, outs: ir(ins[0], ins[1], ins[2], ins[3], ins[4], outs[0], outs[1]),
- dtype=["int64", "float32"],
- name="collect_indices_and_scores",
- tag="collect_indices_and_scores",
- )
-
-
-def _all_class_nms_ir(
- boxes,
- sorted_scores,
- sorted_indices,
- valid_count,
- batch_class,
- num_class,
- num_anchors,
- iou_threshold,
- max_output_size_per_class,
- box_indices,
- selected_scores,
- num_valid_boxes,
- nms_loop,
-):
- ib = tvm.tir.ir_builder.create()
- boxes = ib.buffer_ptr(boxes)
- sorted_scores = ib.buffer_ptr(sorted_scores)
- sorted_indices = ib.buffer_ptr(sorted_indices)
- valid_count = ib.buffer_ptr(valid_count)
- box_indices = ib.buffer_ptr(box_indices)
- num_valid_boxes = ib.buffer_ptr(num_valid_boxes)
-
- if selected_scores is not None:
- selected_scores = ib.buffer_ptr(selected_scores)
-
- if isinstance(iou_threshold, float):
- iou_threshold = tvm.tir.FloatImm("float32", iou_threshold)
-
- if isinstance(max_output_size_per_class, int):
- max_output_size_per_class = tvm.tir.const(max_output_size_per_class)
-
- def calc_overlap(i, j, k):
- offset_j = sorted_indices[i, j] * 4
- offset_k = sorted_indices[i, k] * 4
- batch_id = i // num_class
- base_bbox_idx = batch_id * num_anchors * 4
- return calculate_overlap(
- boxes,
- base_bbox_idx + offset_j,
- base_bbox_idx + offset_k,
- )
-
- def on_new_valid_box(ib, tid, num_current_valid_box, i, j):
- with ib.if_scope(tid + 0 == 0):
- box_indices[i, num_current_valid_box] = sorted_indices[i, j]
-
- if selected_scores is not None:
- selected_scores[i, num_current_valid_box] = sorted_scores[i, j]
-
- def on_new_invalidated_box(*_):
- pass
-
- def needs_bbox_check(*_):
- return tvm.tir.const(True)
-
- return nms_loop(
- ib,
- batch_class,
- tvm.tir.IntImm("int32", -1), # top_k
- iou_threshold,
- max_output_size_per_class,
- valid_count,
- on_new_valid_box,
- on_new_invalidated_box,
- needs_bbox_check,
- calc_overlap,
- sorted_scores,
- num_valid_boxes,
- )
-
-
-def run_all_class_nms(
- boxes,
- sorted_scores,
- sorted_indices,
- valid_count,
- max_output_size_per_class,
- iou_threshold,
- nms_loop,
- return_scores=False,
-):
- """The core all class NMS routine
-
- Parameters
- ----------
- boxes : tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_boxes, 4)
-
- sorted_scores: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes)
- One of the outputs from argsort
-
- sorted_indices: tvm.te.Tensor
- 2-D tensor with shape (batch_size * num_classes, num_boxes)
- The other output from argsort
-
- valid_count: tvm.te.Tensor
- 1-D tensor with shape (batch_size * num_classes,), representing
- the number of boxes whose score is above score_threshold, per batch and class
-
- max_output_boxes_per_class : int or tvm.te.Tensor, optional
- The maxinum number of output selected boxes per class
-
- iou_threshold : float or tvm.te.Tensor, optionaIl
- IoU test threshold
-
- nms_loop : function
- A core NMS loop, see its usage in vision/nms.py and cuda/nms.py
-
- return_scores : bool, optional
- Whether or not to return selected scores, needed by the tensorflow output format.
-
- Returns
- -------
- out : a list of tvm.te.Tensor
- The output is three tensors, the first and second are indices and scores of size
- (batch_size * num_class, num_boxes), and the third is a tensor
- num_selected_boxes of shape (batch_size * num_class,) representing the total number of
- selected boxes per batch and class. If return_scores is False, the second output is
- None.
- """
- batch, num_boxes, _ = boxes.shape
- batch_class = sorted_scores.shape[0]
- num_class = batch_class // batch
-
- if return_scores is False:
- selected_indices, num_detections = te.extern(
- [(batch_class, num_boxes), (1, batch_class)],
- [boxes, sorted_scores, sorted_indices, valid_count],
- lambda ins, outs: _all_class_nms_ir(
- ins[0], # boxes
- ins[1], # sorted_scores
- ins[2], # sorted_indices
- ins[3], # valid_count
- batch_class,
- num_class,
- num_boxes,
- iou_threshold,
- max_output_size_per_class,
- outs[0], # box_indices
- None, # scores
- outs[1], # num_selected_boxes
- nms_loop,
- ),
- dtype=["int32", "int32"],
- name="all_class_nms",
- tag="all_class_nms",
- )
- return selected_indices, None, num_detections
-
- return te.extern(
- [(batch_class, num_boxes), (batch_class, num_boxes), (1, batch_class)],
- [boxes, sorted_scores, sorted_indices, valid_count],
- lambda ins, outs: _all_class_nms_ir(
- ins[0], # boxes
- ins[1], # sorted_scores
- ins[2], # sorted_indices
- ins[3], # valid_count
- batch_class,
- num_class,
- num_boxes,
- iou_threshold,
- max_output_size_per_class,
- outs[0], # box_indices
- outs[1], # selected scores
- outs[2], # num_selected_boxes
- nms_loop,
- ),
- dtype=["int32", "float32", "int32"],
- name="all_class_nms",
- tag="all_class_nms",
- )
diff --git a/python/tvm/topi/vision/rcnn/__init__.py b/python/tvm/topi/vision/rcnn/__init__.py
deleted file mode 100644
index e5693e8694..0000000000
--- a/python/tvm/topi/vision/rcnn/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# 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.
-
-# pylint: disable=wildcard-import
-"""Faster R-CNN and Mask R-CNN operators"""
-from .roi_align import *
-from .roi_pool import *
-from .proposal import *
diff --git a/python/tvm/topi/vision/rcnn/proposal.py b/python/tvm/topi/vision/rcnn/proposal.py
deleted file mode 100644
index 12a0d6bcf0..0000000000
--- a/python/tvm/topi/vision/rcnn/proposal.py
+++ /dev/null
@@ -1,448 +0,0 @@
-# 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.
-# pylint: disable=invalid-name, singleton-comparison, bad-continuation
-"""Proposal operator"""
-import math
-import tvm
-from tvm import te
-from ...utils import get_const_tuple, get_const_int
-from ...sort import argsort
-
-
-def generate_anchor(ratio, scale, base_size):
- """Generate anchor"""
- w = h = float(base_size)
- x_ctr = 0.5 * (w - 1.0)
- y_ctr = 0.5 * (h - 1.0)
- size = w * h
- size_ratios = math.floor(size / ratio)
- new_w = math.floor(math.sqrt(size_ratios) + 0.5) * scale
- new_h = math.floor((new_w / scale * ratio) + 0.5) * scale
- return (
- x_ctr - 0.5 * (new_w - 1.0),
- y_ctr - 0.5 * (new_h - 1.0),
- x_ctr + 0.5 * (new_w - 1.0),
- y_ctr + 0.5 * (new_h - 1.0),
- )
-
-
-def reg_bbox(x1, y1, x2, y2, dx, dy, dw, dh):
- """Bounding box regression function"""
- bbox_w = x2 - x1 + 1.0
- bbox_h = y2 - y1 + 1.0
- ctr_x = x1 + 0.5 * (bbox_w - 1.0)
- ctr_y = y1 + 0.5 * (bbox_h - 1.0)
-
- pred_ctr_x = dx * bbox_w + ctr_x
- pred_ctr_y = dy * bbox_h + ctr_y
- pred_w = te.exp(dw) * bbox_w
- pred_h = te.exp(dh) * bbox_h
-
- pred_x1 = pred_ctr_x - 0.5 * (pred_w - 1.0)
- pred_y1 = pred_ctr_y - 0.5 * (pred_h - 1.0)
- pred_x2 = pred_ctr_x + 0.5 * (pred_w - 1.0)
- pred_y2 = pred_ctr_y + 0.5 * (pred_h - 1.0)
- return pred_x1, pred_y1, pred_x2, pred_y2
-
-
-def reg_iou(x1, y1, x2, y2, dx1, dy1, dx2, dy2):
- """Bounding box regression function"""
- pred_x1 = x1 + dx1
- pred_y1 = y1 + dy1
- pred_x2 = x2 + dx2
- pred_y2 = y2 + dy2
- return pred_x1, pred_y1, pred_x2, pred_y2
-
-
-def predict_bbox_ir(
- cls_prob_buf,
- bbox_pred_buf,
- im_info_buf,
- out_buf,
- scales,
- ratios,
- feature_stride,
- rpn_min_size,
- iou_loss,
-):
- """Predict bounding boxes based on anchors, scores and deltas.
-
- Parameters
- ----------
- cls_prob_buf : tvm.te.schedule.Buffer
- 4-D with shape [batch, 2 * num_anchors, height, width]
-
- bbox_pred_buf : tvm.te.schedule.Buffer
- 4-D with shape [batch, 4 * num_anchors, height, width]
-
- im_info_buf : tvm.te.schedule.Buffer
- 2-D with shape [batch, 3]
-
- out_buf : tvm.te.schedule.Buffer
- 3-D with shape [batch, num_bbox, 5]
- The last dimension is in format of [w_start, h_start, w_end, h_end, score]
-
- scales : list/tuple of float
- Scales of anchor windows.
-
- ratios : list/tuple of float
- Ratios of anchor windows.
-
- feature_stride : int
- The size of the receptive field each unit in the convolution layer of the rpn, for example
- the product of all stride's prior to this layer.
-
- rpn_min_size : int
- Minimum height or width in proposal.
-
- iou_loss : bool
- Usage of IoU loss.
-
- Returns
- -------
- stmt : Stmt
- The result IR statement.
- """
- batch, num_anchors, height, width = get_const_tuple(cls_prob_buf.shape)
- num_anchors //= 2
- ib = tvm.tir.ir_builder.create()
-
- p_score = ib.buffer_ptr(cls_prob_buf)
- p_delta = ib.buffer_ptr(bbox_pred_buf)
- p_im_info = ib.buffer_ptr(im_info_buf)
- p_out = ib.buffer_ptr(out_buf)
-
- idxm = tvm.tir.indexmod
- idxd = tvm.tir.indexdiv
-
- with ib.for_range(0, batch * height * width) as tid:
- w = idxm(tid, width)
- h = idxm(idxd(tid, width), height)
- b = idxd(idxd(tid, width), height)
-
- for k in range(num_anchors):
- out_index = tid * num_anchors + k
- ratio = ratios[k // len(scales)]
- scale = scales[k % len(scales)]
- anchor = generate_anchor(ratio, scale, feature_stride)
- im_height = p_im_info[b * 3]
- im_width = p_im_info[b * 3 + 1]
- x1 = anchor[0] + w * feature_stride
- y1 = anchor[1] + h * feature_stride
- x2 = anchor[2] + w * feature_stride
- y2 = anchor[3] + h * feature_stride
-
- delta = [
- p_delta[((((b * num_anchors + k) * 4 + i) * height + h) * width + w)]
- for i in range(4)
- ]
- regression_func = reg_iou if iou_loss else reg_bbox
- pred_x1, pred_y1, pred_x2, pred_y2 = regression_func(x1, y1, x2, y2, *delta)
-
- pred_x1 = tvm.te.max(tvm.te.min(pred_x1, im_width - 1.0), 0.0)
- pred_y1 = tvm.te.max(tvm.te.min(pred_y1, im_height - 1.0), 0.0)
- pred_x2 = tvm.te.max(tvm.te.min(pred_x2, im_width - 1.0), 0.0)
- pred_y2 = tvm.te.max(tvm.te.min(pred_y2, im_height - 1.0), 0.0)
-
- real_height = (im_height / feature_stride).astype("int32")
- real_width = (im_width / feature_stride).astype("int32")
-
- bbox_w = pred_x2 - pred_x1 + 1.0
- bbox_h = pred_y2 - pred_y1 + 1.0
- min_size = p_im_info[b * 3 + 2] * rpn_min_size
-
- pred_score = p_score[((b * num_anchors * 2 + num_anchors + k) * height + h) * width + w]
- pred_score = tvm.tir.Select(
- tvm.tir.any(h >= real_height, w >= real_width), -1.0, pred_score
- )
- p_out[out_index * 5 + 0] = pred_x1
- p_out[out_index * 5 + 1] = pred_y1
- p_out[out_index * 5 + 2] = pred_x2
- p_out[out_index * 5 + 3] = pred_y2
- p_out[out_index * 5 + 4] = pred_score
-
- with ib.if_scope(tvm.tir.any(bbox_w < min_size, bbox_h < min_size)):
- p_out[out_index * 5 + 0] -= min_size / 2.0
- p_out[out_index * 5 + 1] -= min_size / 2.0
- p_out[out_index * 5 + 2] += min_size / 2.0
- p_out[out_index * 5 + 3] += min_size / 2.0
- p_out[out_index * 5 + 4] = -1.0
-
- return ib.get()
-
-
-def argsort_ir(data_buf, out_index_buf):
- """Batched odd-even transposition sort.
-
- Parameters
- ----------
- data_buf : tvm.te.schedule.Buffer
- 2-D with shape [batch, num_bbox]
-
- out_index_buf : tvm.te.schedule.Buffer
- 2-D with shape [batch, num_bbox]. Indices of data in sorted order.
-
- Returns
- -------
- stmt : Stmt
- The result IR statement.
- """
- batch, num_bbox = get_const_tuple(data_buf.shape)
- ib = tvm.tir.ir_builder.create()
- p_data = ib.buffer_ptr(data_buf)
- index_out = ib.buffer_ptr(out_index_buf)
- temp_data = ib.allocate("float32", (1,), name="temp_data", scope="local")
- temp_index = ib.allocate("int32", (1,), name="temp_index", scope="local")
- idxm = tvm.tir.indexmod
- with ib.for_range(0, batch, kind="unroll") as b:
- start = b * num_bbox
- for i in range(2):
- with ib.for_range(0, (num_bbox + 1) // 2) as tid:
- bbox_id = tid * 2 + i
- with ib.if_scope(bbox_id < num_bbox):
- index_out[start + bbox_id] = bbox_id
- with ib.for_range(0, num_bbox) as k:
- with ib.for_range(0, (num_bbox + 1) // 2) as tid:
- offset = start + 2 * tid + idxm(k, 2)
- with ib.if_scope(
- tvm.tir.all(offset + 1 < num_bbox, p_data[offset] < p_data[offset + 1])
- ):
- temp_data[0] = p_data[offset]
- p_data[offset] = p_data[offset + 1]
- p_data[offset + 1] = temp_data[0]
- temp_index[0] = index_out[offset]
- index_out[offset] = index_out[offset + 1]
- index_out[offset + 1] = temp_index[0]
- return ib.get()
-
-
-def nms_ir(sorted_bbox_buf, out_buf, nms_threshold):
- """Non-maximum suppression.
-
- Parameters
- ----------
- sorted_bbox_buf : tvm.te.schedule.Buffer
- 3-D with shape [batch, num_bbox, 5]. The last dimension is in format of
- [w_start, h_start, w_end, h_end, score].
-
- out_buf : tvm.te.schedule.Buffer
- 2-D with shape [batch, num_bbox]. Boolean mask of whether a bounding box should be removed.
-
- nms_threshold : float
- Non-maximum suppression threshold.
-
- Returns
- -------
- stmt : Stmt
- The result IR statement.
- """
-
- def calculate_overlap(out_tensor, box_a_idx, box_b_idx):
- """Calculate overlap of two boxes."""
- w = tvm.te.max(
- 0.0,
- tvm.te.min(out_tensor[box_a_idx + 2], out_tensor[box_b_idx + 2])
- - tvm.te.max(out_tensor[box_a_idx], out_tensor[box_b_idx])
- + 1.0,
- )
- h = tvm.te.max(
- 0.0,
- tvm.te.min(out_tensor[box_a_idx + 3], out_tensor[box_b_idx + 3])
- - tvm.te.max(out_tensor[box_a_idx + 1], out_tensor[box_b_idx + 1])
- + 1.0,
- )
- i = w * h
- u = (
- (out_tensor[box_a_idx + 2] - out_tensor[box_a_idx] + 1.0)
- * (out_tensor[box_a_idx + 3] - out_tensor[box_a_idx + 1] + 1.0)
- + (out_tensor[box_b_idx + 2] - out_tensor[box_b_idx] + 1.0)
- * (out_tensor[box_b_idx + 3] - out_tensor[box_b_idx + 1] + 1.0)
- - i
- )
- return i / u
-
- batch, num_bbox = get_const_tuple(out_buf.shape)
- ib = tvm.tir.ir_builder.create()
- p_data = ib.buffer_ptr(sorted_bbox_buf)
- p_out = ib.buffer_ptr(out_buf)
- with ib.for_range(0, batch, kind="unroll", name="n") as b:
- base_idx = b * num_bbox
- for i in range(num_bbox):
- p_out[base_idx + i] = False
- with ib.for_range(0, num_bbox - 1) as l:
- with ib.for_range(0, num_bbox) as i:
- with ib.if_scope(tvm.tir.all(i < num_bbox, i > l, p_out[base_idx + l] == False)):
- iou = calculate_overlap(p_data, (base_idx + l) * 5, (base_idx + i) * 5)
- with ib.if_scope(iou > nms_threshold):
- p_out[base_idx + i] = True
- return ib.get()
-
-
-def prepare_output_ir(sorted_bbox_buf, remove_mask_buf, out_buf):
- """Copy output after applying nms to continuous memory.
-
- Parameters
- ----------
- sorted_bbox_buf : tvm.te.schedule.Buffer
- 3-D with shape [batch, num_bbox, 5]. The last dimension is in format of
- [w_start, h_start, w_end, h_end, score].
-
- remove_mask_buf : tvm.te.schedule.Buffer
- 2-D with shape [batch, num_bbox]. Boolean mask of whether a bounding box should be removed.
-
- out_buf : tvm.te.schedule.Buffer
- 2-D with shape [batch * rpn_post_nms_top_n, 5]. The last dimension is in format of
- [batch_index, w_start, h_start, w_end, h_end].
-
- Returns
- -------
- stmt : Stmt
- The result IR statement.
- """
- batch, num_bbox, _ = get_const_tuple(sorted_bbox_buf.shape)
- rpn_post_nms_top_n = get_const_int(out_buf.shape[0]) // batch
- ib = tvm.tir.ir_builder.create()
- i = ib.allocate("int32", (batch,), "i", scope="local")
- p_sorted_bbox = ib.buffer_ptr(sorted_bbox_buf)
- p_remove = ib.buffer_ptr(remove_mask_buf)
- p_out = ib.buffer_ptr(out_buf)
-
- nkeep = ib.allocate("int32", (batch,), "nkeep", scope="local")
-
- with ib.for_range(0, batch) as b:
- nkeep[b] = 0
- i[b] = 0
-
- with ib.for_range(0, num_bbox) as j:
- with ib.for_range(0, batch) as b:
- with ib.if_scope(p_remove[b * num_bbox + j] == False):
- nkeep[b] += 1
- with ib.for_range(0, batch) as b:
- with ib.if_scope(nkeep[b] > 0):
- with ib.for_range(
- 0, te.ceil(tvm.tir.const(rpn_post_nms_top_n, "float32") / nkeep[b]).astype("int32")
- ):
- with ib.for_range(0, num_bbox) as j:
- offset_j = (b * num_bbox + j) * 5
- offset_i = (b * rpn_post_nms_top_n + i[b]) * 5
- with ib.if_scope(
- tvm.tir.all(
- i[b] < rpn_post_nms_top_n, p_remove[(b * num_bbox + j)] == False
- )
- ):
- p_out[offset_i] = tvm.tir.Cast("float32", b)
- with ib.for_range(0, 4, kind="unroll") as k:
- p_out[offset_i + k + 1] = p_sorted_bbox[offset_j + k]
- i[b] = i[b] + 1
-
- body = ib.get()
- return body
-
-
-def proposal(
- cls_prob,
- bbox_pred,
- im_info,
- scales,
- ratios,
- feature_stride,
- threshold,
- rpn_pre_nms_top_n,
- rpn_post_nms_top_n,
- rpn_min_size,
- iou_loss,
-):
- """Proposal operator.
-
- Parameters
- ----------
- cls_prob : tvm.te.Tensor
- 4-D with shape [batch, 2 * num_anchors, height, width]
-
- bbox_pred : tvm.te.Tensor
- 4-D with shape [batch, 4 * num_anchors, height, width]
-
- im_info : tvm.te.Tensor
- 2-D with shape [batch, 3]
-
- scales : list/tuple of float
- Scales of anchor windows.
-
- ratios : list/tuple of float
- Ratios of anchor windows.
-
- feature_stride : int
- The size of the receptive field each unit in the convolution layer of the rpn, for example
- the product of all stride's prior to this layer.
-
- threshold : float
- Non-maximum suppression threshold.
-
- rpn_pre_nms_top_n : int
- Number of top scoring boxes to apply NMS. -1 to use all boxes.
-
- rpn_post_nms_top_n : int
- Number of top scoring boxes to keep after applying NMS to RPN proposals.
-
- rpn_min_size : int
- Minimum height or width in proposal.
-
- iou_loss : bool
- Usage of IoU loss.
-
- Returns
- -------
- out : tvm.te.Tensor
- 2-D tensor with shape [batch * rpn_post_nms_top_n, 5]. The last dimension is in format of
- [batch_index, w_start, h_start, w_end, h_end].
- """
- # pylint: disable=unused-argument
- batch, _, height, width = get_const_tuple(cls_prob.shape)
- num_anchors = len(scales) * len(ratios)
- num_bbox = height * width * num_anchors
- rpn_pre_nms_top_n = min(rpn_pre_nms_top_n, num_bbox) if rpn_pre_nms_top_n > 0 else num_bbox
-
- bbox = te.extern(
- (batch, num_bbox, 5),
- [cls_prob, bbox_pred, im_info],
- lambda ins, outs: predict_bbox_ir(
- ins[0], ins[1], ins[2], outs[0], scales, ratios, feature_stride, rpn_min_size, iou_loss
- ),
- dtype=bbox_pred.dtype,
- )
- score = te.compute((batch, num_bbox), lambda b, i: bbox[b, i, 4], tag="bbox_score")
- valid_count_shape = (1,)
- valid_count = te.compute(valid_count_shape, lambda i: num_bbox)
- sorted_index = argsort(score, valid_count=valid_count, axis=1, is_ascend=False)
- sorted_bbox = te.compute(
- (batch, rpn_pre_nms_top_n, 5),
- lambda b, i, j: bbox[b, sorted_index[b, i], j],
- tag="sorted_bbox",
- )
- nms_remove_mask = te.extern(
- (batch, rpn_pre_nms_top_n),
- [sorted_bbox],
- lambda ins, outs: nms_ir(ins[0], outs[0], threshold),
- dtype="bool",
- )
- nms_out = te.extern(
- (batch * rpn_post_nms_top_n, 5),
- [sorted_bbox, nms_remove_mask],
- lambda ins, outs: prepare_output_ir(ins[0], ins[1], outs[0]),
- dtype=sorted_bbox.dtype,
- )
- return nms_out
diff --git a/python/tvm/topi/vision/rcnn/roi_align.py b/python/tvm/topi/vision/rcnn/roi_align.py
deleted file mode 100644
index 238e029643..0000000000
--- a/python/tvm/topi/vision/rcnn/roi_align.py
+++ /dev/null
@@ -1,228 +0,0 @@
-# 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.
-# pylint: disable=invalid-name
-"""Roi align operator"""
-import tvm
-from tvm import te
-from ...utils import get_const_tuple
-from ...cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
-
-
-def _sample_common(
- i,
- c,
- ph,
- pw,
- rois,
- pooled_size_h,
- pooled_size_w,
- spatial_scale,
- sample_ratio,
- dtype,
- avg_mode,
- bilinear_func,
-):
- roi = rois[i]
- batch_index = roi[0].astype("int32")
- roi_start_w, roi_start_h, roi_end_w, roi_end_h = roi[1], roi[2], roi[3], roi[4]
- roi_start_h *= spatial_scale
- roi_end_h *= spatial_scale
- roi_start_w *= spatial_scale
- roi_end_w *= spatial_scale
-
- # force malformed ROIs to be 1x1
- roi_h = tvm.te.max(roi_end_h - roi_start_h, tvm.tir.const(1.0, dtype))
- roi_w = tvm.te.max(roi_end_w - roi_start_w, tvm.tir.const(1.0, dtype))
-
- bin_h = roi_h / pooled_size_h
- bin_w = roi_w / pooled_size_w
-
- if sample_ratio > 0:
- roi_bin_grid_h = roi_bin_grid_w = tvm.tir.const(sample_ratio, "int32")
- else:
- roi_bin_grid_h = te.ceil(roi_h / pooled_size_h).astype("int32")
- roi_bin_grid_w = te.ceil(roi_w / pooled_size_w).astype("int32")
-
- count = roi_bin_grid_h * roi_bin_grid_w
- rh = te.reduce_axis((0, roi_bin_grid_h), name="rh")
- rw = te.reduce_axis((0, roi_bin_grid_w), name="rw")
- roi_start_h += ph * bin_h
- roi_start_w += pw * bin_w
-
- if avg_mode:
- return te.sum(
- bilinear_func(
- batch_index,
- c,
- roi_start_h + (rh + 0.5) * bin_h / roi_bin_grid_h,
- roi_start_w + (rw + 0.5) * bin_w / roi_bin_grid_w,
- )
- / count,
- axis=[rh, rw],
- )
- # max mode
- return te.max(
- bilinear_func(
- batch_index,
- c,
- roi_start_h + (rh + 0.5) * bin_h / roi_bin_grid_h,
- roi_start_w + (rw + 0.5) * bin_w / roi_bin_grid_w,
- ),
- axis=[rh, rw],
- )
-
-
-def roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1):
- """ROI align operator in NCHW layout.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- 4-D with shape [batch, channel, height, width]
-
- rois : tvm.te.Tensor
- 2-D with shape [num_roi, 5]. The last dimension should be in format of
- [batch_index, w_start, h_start, w_end, h_end]
-
- pooled_size : int or list/tuple of two ints
- output size, or [out_height, out_width]
-
- spatial_scale : float
- Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal
- of total stride in convolutional layers, which should be in range (0.0, 1.0]
-
- mode : int or str
- There are two modes, average and max. For the average mode, you can pass b'avg' or 0, and
- for the max mode, you can pass b'max' or 1.
-
- sample_ratio : int
- Optional sampling ratio of ROI align, using adaptive size by default.
-
- Returns
- -------
- output : tvm.te.Tensor
- 4-D with shape [num_roi, channel, pooled_size, pooled_size]
- """
- avg_mode = mode in (b"avg", 0)
- max_mode = mode in (b"max", 1)
- assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode."
- dtype = rois.dtype
- _, channel, height, width = get_const_tuple(data.shape)
- num_roi, _ = get_const_tuple(rois.shape)
-
- if isinstance(pooled_size, int):
- pooled_size_h = pooled_size_w = pooled_size
- else:
- pooled_size_h, pooled_size_w = pooled_size
-
- def _bilinear(i, c, y, x):
- outside = tvm.tir.any(y < -1.0, x < -1.0, y > height, x > width)
- y = tvm.te.min(tvm.te.max(y, 0.0), height - 1)
- x = tvm.te.min(tvm.te.max(x, 0.0), width - 1)
- val = bilinear_sample_nchw(data, (i, c, y, x), height - 1, width - 1)
- return tvm.tir.if_then_else(outside, 0.0, val)
-
- def _sample(i, c, ph, pw):
- return _sample_common(
- i,
- c,
- ph,
- pw,
- rois,
- pooled_size_h,
- pooled_size_w,
- spatial_scale,
- sample_ratio,
- dtype,
- avg_mode,
- _bilinear,
- )
-
- return te.compute(
- (num_roi, channel, pooled_size_h, pooled_size_w), _sample, tag="pool,roi_align_nchw"
- )
-
-
-def roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1):
- """ROI align operator in NHWC layout.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- 4-D with shape [batch, height, width, channel]
-
- rois : tvm.te.Tensor
- 2-D with shape [num_roi, 5]. The last dimension should be in format of
- [batch_index, w_start, h_start, w_end, h_end]
-
- pooled_size : int or list/tuple of two ints
- output size, or [out_height, out_width]
-
- spatial_scale : float
- Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal
- of total stride in convolutional layers, which should be in range (0.0, 1.0]
-
- mode : int or str
- There are two modes, average and max. For the average mode, you can pass b'avg' or 0, and
- for the max mode, you can pass b'max' or 1.
-
- sample_ratio : int
- Optional sampling ratio of ROI align, using adaptive size by default.
-
- Returns
- -------
- output : tvm.te.Tensor
- 4-D with shape [num_roi, pooled_size, pooled_size, channel]
- """
- avg_mode = mode in (b"avg", 0)
- max_mode = mode in (b"max", 1)
- assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode."
- dtype = rois.dtype
- _, height, width, channel = get_const_tuple(data.shape)
- num_roi, _ = get_const_tuple(rois.shape)
-
- if isinstance(pooled_size, int):
- pooled_size_h = pooled_size_w = pooled_size
- else:
- pooled_size_h, pooled_size_w = pooled_size
-
- def _bilinear(i, c, y, x):
- outside = tvm.tir.any(y < -1.0, x < -1.0, y > height, x > width)
- y = tvm.te.min(tvm.te.max(y, 0.0), height - 1)
- x = tvm.te.min(tvm.te.max(x, 0.0), width - 1)
- val = bilinear_sample_nhwc(data, (i, y, x, c), height - 1, width - 1)
- return tvm.tir.if_then_else(outside, 0.0, val)
-
- def _sample(i, ph, pw, c):
- return _sample_common(
- i,
- c,
- ph,
- pw,
- rois,
- pooled_size_h,
- pooled_size_w,
- spatial_scale,
- sample_ratio,
- dtype,
- avg_mode,
- _bilinear,
- )
-
- return te.compute(
- (num_roi, pooled_size_h, pooled_size_w, channel), _sample, tag="pool,roi_align_nchw"
- )
diff --git a/python/tvm/topi/vision/rcnn/roi_pool.py b/python/tvm/topi/vision/rcnn/roi_pool.py
deleted file mode 100644
index dd1429bcb3..0000000000
--- a/python/tvm/topi/vision/rcnn/roi_pool.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# 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.
-# pylint: disable=invalid-name
-"""ROI pool operator"""
-import tvm
-from tvm import te
-from ...utils import get_const_tuple
-
-
-def roi_pool_nchw(data, rois, pooled_size, spatial_scale):
- """ROI pool operator in NCHW layout.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- 4-D with shape [batch, channel, height, width]
-
- rois : tvm.te.Tensor
- 2-D with shape [num_roi, 5]. The last dimension should be in format of
- [batch_index, w_start, h_start, w_end, h_end]
-
- pooled_size : int or list/tuple of two ints
- output size, or [out_height, out_width]
-
- spatial_scale : float
- Ratio of input feature map height (or w) to raw image height (or w). Equals the reciprocal
- of total stride in convolutional layers, which should be in range (0.0, 1.0]
-
- Returns
- -------
- output : tvm.te.Tensor
- 4-D with shape [num_roi, channel, pooled_size, pooled_size]
- """
- dtype = rois.dtype
- _, channel, height, width = get_const_tuple(data.shape)
- num_roi, _ = get_const_tuple(rois.shape)
-
- if isinstance(pooled_size, int):
- pooled_size_h = pooled_size_w = pooled_size
- else:
- pooled_size_h, pooled_size_w = pooled_size
-
- def _pool(i, c, ph, pw):
- roi = rois[i]
- batch_index = roi[0].astype("int32")
- roi_start_w, roi_start_h, roi_end_w, roi_end_h = roi[1], roi[2], roi[3], roi[4]
-
- roi_start_h = te.round(roi_start_h * spatial_scale).astype("int32")
- roi_start_w = te.round(roi_start_w * spatial_scale).astype("int32")
- roi_end_h = te.round(roi_end_h * spatial_scale).astype("int32")
- roi_end_w = te.round(roi_end_w * spatial_scale).astype("int32")
-
- # force malformed ROIs to be 1x1
- roi_h = tvm.te.max(roi_end_h - roi_start_h + 1, tvm.tir.const(1, "int32"))
- roi_w = tvm.te.max(roi_end_w - roi_start_w + 1, tvm.tir.const(1, "int32"))
-
- bin_h = roi_h.astype(dtype) / pooled_size_h
- bin_w = roi_w.astype(dtype) / pooled_size_w
-
- # use epsilon to prevent floating point precision loss in floor/ceil
- epsilon = tvm.tir.const(0.00001, dtype)
- hstart = te.floor(ph * bin_h + epsilon).astype("int32")
- wstart = te.floor(pw * bin_w + epsilon).astype("int32")
- hend = te.ceil((ph + 1) * bin_h - epsilon).astype("int32")
- wend = te.ceil((pw + 1) * bin_w - epsilon).astype("int32")
- hstart = tvm.te.min(tvm.te.max(hstart + roi_start_h, 0), height)
- wstart = tvm.te.min(tvm.te.max(wstart + roi_start_w, 0), width)
- hend = tvm.te.min(tvm.te.max(hend + roi_start_h, 0), height)
- wend = tvm.te.min(tvm.te.max(wend + roi_start_w, 0), width)
-
- non_empty = tvm.tir.all(hstart < hend, wstart < wend)
- min_value = lambda dtype: tvm.tir.if_then_else(
- non_empty, tvm.te.min_value(dtype), tvm.tir.const(0.0, dtype)
- )
- # pylint: disable=unnecessary-lambda
- _max = te.comm_reducer(lambda x, y: tvm.te.max(x, y), min_value, name="max")
- rh = te.reduce_axis((0, hend - hstart), "rh")
- rw = te.reduce_axis((0, wend - wstart), "rw")
- return _max(data[batch_index, c, hstart + rh, wstart + rw], axis=[rh, rw])
-
- return te.compute((num_roi, channel, pooled_size_h, pooled_size_w), _pool, tag="pool,roi_pool")
diff --git a/python/tvm/topi/vision/reorg.py b/python/tvm/topi/vision/reorg.py
deleted file mode 100644
index 9883085f9f..0000000000
--- a/python/tvm/topi/vision/reorg.py
+++ /dev/null
@@ -1,42 +0,0 @@
-# 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.
-"""
-REORG Operator
-====================
-Reorg operator, used in darknet.
-"""
-from __future__ import absolute_import as _abs
-from .. import cpp
-
-
-def reorg(data, stride):
- """Reorg forward operators.
-
- Parameters
- ----------
- Input : tvm.te.Tensor
- 4-D with shape [batch, in_channel, in_height, in_width]
-
- stride : int
- Stride value for reorganization
-
- Returns
- -------
- Output : tvm.te.Tensor
- 4-D with shape [batch, out_channel, out_height, out_width]
- """
- return cpp.vision.reorg(data, stride)
diff --git a/python/tvm/topi/vision/ssd/__init__.py b/python/tvm/topi/vision/ssd/__init__.py
deleted file mode 100644
index 1ac388da9a..0000000000
--- a/python/tvm/topi/vision/ssd/__init__.py
+++ /dev/null
@@ -1,22 +0,0 @@
-# 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.
-
-# pylint: disable=wildcard-import
-"""VISION network operators"""
-from __future__ import absolute_import as _abs
-
-from .multibox import *
diff --git a/python/tvm/topi/vision/ssd/multibox.py b/python/tvm/topi/vision/ssd/multibox.py
deleted file mode 100644
index 234bfd7953..0000000000
--- a/python/tvm/topi/vision/ssd/multibox.py
+++ /dev/null
@@ -1,369 +0,0 @@
-# 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.
-# pylint: disable=invalid-name, no-member, too-many-locals, too-many-arguments, undefined-variable
-"""SSD multibox operators"""
-import tvm
-
-from tvm.te import hybrid
-from tvm.tir import exp, sqrt
-
-from tvm import topi
-
-from ..nms import non_max_suppression
-
-
-@hybrid.script
-def hybrid_multibox_prior(data, sizes, ratios, steps, offsets):
- """Hybrid routing for multibox_prior operator.
-
- Parameters
- ----------
- data : tvm.te.Tensor or numpy NDArray
- 4-D tensor with shape [batch, channel, height, width]]
-
- sizes : tvm ConsExpr
- Sizes for anchor boxes.
-
- ratios : tvm ConsExpr
- Ratios for anchor boxes.
-
- steps : tvm ConsExpr
- Priorbox step across y and x, -1 for auto calculation.
-
- offsets : tvm ConsExpr
- Priorbox center offsets, y and x respectively.
-
- Returns
- -------
- output : tvm.te.Tensor or numpy NDArray
- 3-D tensor with shape [1, h_in * w_in * (num_sizes + num_ratios - 1), 4]
- """
- in_height = data.shape[2]
- in_width = data.shape[3]
- num_sizes = len(sizes)
- num_ratios = len(ratios)
- num_boxes = in_height * in_width * (num_sizes + num_ratios - 1)
- output = output_tensor((1, num_boxes, 4), "float32")
- steps_h = steps[0] * 1.0 if steps[0] > 0 else 1.0 / in_height
- steps_w = steps[1] * 1.0 if steps[1] > 0 else 1.0 / in_width
- offset_h = offsets[0]
- offset_w = offsets[1]
-
- # Need to define var out of const_range + if
- w = 0.0
- h = 0.0
-
- for i in parallel(in_height):
- center_h = (i + offset_h) * steps_h
- for j in range(in_width):
- center_w = (j + offset_w) * steps_w
- for k in const_range(num_sizes + num_ratios - 1):
- if k < num_sizes:
- w = float32(sizes[k] * in_height) / in_width / 2.0
- h = sizes[k] / 2.0
- else:
- w = (
- float32(sizes[0] * in_height)
- / in_width
- * sqrt(ratios[k - num_sizes + 1] * 1.0)
- / 2.0
- )
- h = sizes[0] / sqrt(ratios[k - num_sizes + 1] * 1.0) / 2.0
- count = (
- i * in_width * (num_sizes + num_ratios - 1)
- + j * (num_sizes + num_ratios - 1)
- + k
- )
- output[0, count, 0] = center_w - w
- output[0, count, 1] = center_h - h
- output[0, count, 2] = center_w + w
- output[0, count, 3] = center_h + h
-
- return output
-
-
-def multibox_prior(data, sizes=(1,), ratios=(1,), steps=(-1, -1), offsets=(0.5, 0.5), clip=False):
- """Generate prior(anchor) boxes from data, sizes and ratios.
-
- Parameters
- ----------
- data : tvm.te.Tensor
- 4-D with shape [batch, c_in, h_in, w_in]]
-
- sizes : tuple of float
- Tuple of sizes for anchor boxes.
-
- ratios : tuple of float
- Tuple of ratios for anchor boxes.
-
- steps : Tuple of float
- Priorbox step across y and x, -1 for auto calculation.
-
- offsets : tuple of int
- Priorbox center offsets, y and x respectively.
-
- clip : boolean
- Whether to clip out-of-boundary boxes.
-
- Returns
- -------
- out : tvm.te.Tensor
- 3-D tensor with shape [1, h_in * w_in * (num_sizes + num_ratios - 1), 4]
- """
- out = hybrid_multibox_prior(
- data,
- tvm.runtime.convert(sizes),
- tvm.runtime.convert(ratios),
- tvm.runtime.convert(steps),
- tvm.runtime.convert(offsets),
- )
- if clip:
- out = topi.clip(out, 0, 1)
- return out
-
-
-@hybrid.script
-def _hybrid_transform_loc(anchor, pred_loc, variance, clip, batch_idx, anchor_idx):
- """Transform prior anchor box to output box through location predictions."""
- al = anchor[0, anchor_idx, 0]
- at = anchor[0, anchor_idx, 1]
- ar = anchor[0, anchor_idx, 2]
- ab = anchor[0, anchor_idx, 3]
-
- px = pred_loc[batch_idx, 0]
- py = pred_loc[batch_idx, 1]
- pw = pred_loc[batch_idx, 2]
- ph = pred_loc[batch_idx, 3]
-
- vx = variance[0]
- vy = variance[1]
- vw = variance[2]
- vh = variance[3]
-
- output = output_tensor((4,), pred_loc.dtype)
-
- aw = ar - al
- ah = ab - at
- ax = (al + ar) / 2.0
- ay = (at + ab) / 2.0
- ox = px * vx * aw + ax
- oy = py * vy * ah + ay
- ow = exp(pw * vw) * aw / 2.0
- oh = exp(ph * vh) * ah / 2.0
- output[0] = max(0.0, min(1.0, ox - ow)) if clip else ox - ow
- output[1] = max(0.0, min(1.0, oy - oh)) if clip else oy - oh
- output[2] = max(0.0, min(1.0, ox + ow)) if clip else ox + ow
- output[3] = max(0.0, min(1.0, oy + oh)) if clip else oy + oh
- return output
-
-
-@hybrid.script
-def hybrid_multibox_transform_loc(
- cls_prob,
- loc_pred,
- anchor,
- clip,
- threshold,
- variances,
- keep_background,
-):
- """Hybrid routing for transform location in multibox_detection operator.
-
- Parameters
- ----------
- cls_prob : tvm.te.Tensor or numpy NDArray
- 3-D tensor of class probabilities.
-
- loc_pred : tvm.te.Tensor or numpy NDArray
- 2-D tensor of location regression predictions.
-
- anchor : tvm.te.Tensor or numpy NDArray
- 3-D tensor of prior anchor boxes.
-
- clip : tvm.tir.const
- Whether to clip out-of-boundary boxes.
-
- threshold : tvm.tir.const
- Threshold to be a positive prediction.
-
- variances : tvm.nd.NDArray
- Variances to be decoded from box regression output.
-
- keep_background : tvm.tir.const
- Whether to keep boxes detected as background or not.
-
- Returns
- -------
- out_loc : tvm.te.Tensor or numpy NDArray
- 3-D tensor of transformed location.
-
- valid_count : tvm.te.Tensor or numpy NDArray
- 1_d tensor of valid counts for boxes.
- """
- batch_size = cls_prob.shape[0]
- num_classes = cls_prob.shape[1]
- num_anchors = cls_prob.shape[2]
- pred_coord = allocate(
- (
- batch_size,
- 4,
- ),
- loc_pred.dtype,
- )
- out_loc = output_tensor((batch_size, num_anchors, 6), loc_pred.dtype)
- valid_count = output_tensor((batch_size,), "int32")
-
- start_cls_idx = 0 if keep_background else 1
-
- for i in parallel(batch_size):
- valid_count[i] = 0
- for j in range(num_anchors):
- # Find the predicted class id and probability
- score = -1.0
- cls_id = 0
- for k in range(start_cls_idx, num_classes):
- temp = cls_prob[i, k, j]
- cls_id = k if temp > score else cls_id
- score = max(temp, score)
- if cls_id > 0 and score < threshold:
- cls_id = 0
- # [id, prob, xmin, ymin, xmax, ymax]
- # Remove background if 'keep_background=False', restore original id
- if keep_background or cls_id > 0:
- out_loc[i, valid_count[i], 0] = cls_id - 0.0 if keep_background else cls_id - 1.0
- out_loc[i, valid_count[i], 1] = score
- for l in range(4):
- pred_coord[i, l] = loc_pred[i, j * 4 + l]
- out_coord = _hybrid_transform_loc(anchor, pred_coord, variances, clip, i, j)
- out_loc[i, valid_count[i], 2] = out_coord[0]
- out_loc[i, valid_count[i], 3] = out_coord[1]
- out_loc[i, valid_count[i], 4] = out_coord[2]
- out_loc[i, valid_count[i], 5] = out_coord[3]
- valid_count[i] += 1
-
- return out_loc, valid_count
-
-
-def multibox_transform_loc(
- cls_prob,
- loc_pred,
- anchor,
- clip=True,
- threshold=0.01,
- variances=(0.1, 0.1, 0.2, 0.2),
- keep_background=False,
-):
- """Location transformation for multibox detection
-
- Parameters
- ----------
- cls_prob : tvm.te.Tensor
- Class probabilities.
-
- loc_pred : tvm.te.Tensor
- Location regression predictions.
-
- anchor : tvm.te.Tensor
- Prior anchor boxes.
-
- clip : boolean
- Whether to clip out-of-boundary boxes.
-
- threshold : float
- Threshold to be a positive prediction.
-
- variances : tuple of float
- Variances to be decoded from box regression output.
-
- keep_background : boolean
- Whether to keep boxes detected as background or not.
-
- Returns
- -------
- ret : tuple of tvm.te.Tensor
- """
-
- return hybrid_multibox_transform_loc(
- cls_prob,
- loc_pred,
- anchor,
- tvm.tir.const(clip, "bool"),
- tvm.tir.const(threshold, "float32"),
- tvm.runtime.convert(variances),
- tvm.tir.const(keep_background, "bool"),
- )
-
-
-def multibox_detection(
- cls_prob,
- loc_pred,
- anchor,
- clip=True,
- threshold=0.01,
- nms_threshold=0.5,
- force_suppress=False,
- variances=(0.1, 0.1, 0.2, 0.2),
- nms_topk=-1,
-):
- """Convert multibox detection predictions.
-
- Parameters
- ----------
- cls_prob : tvm.te.Tensor
- Class probabilities.
-
- loc_pred : tvm.te.Tensor
- Location regression predictions.
-
- anchor : tvm.te.Tensor
- Prior anchor boxes.
-
- clip : boolean
- Whether to clip out-of-boundary boxes.
-
- nms_threshold : float
- Non-maximum suppression threshold.
-
- force_suppress : boolean
- Whether to suppress all detections regardless of class_id.
-
- threshold : float
- Threshold to be a positive prediction.
-
- variances : tuple of float
- Variances to be decoded from box regression output.
-
- nms_topk : int
- Keep maximum top k detections before nms, -1 for no limit.
-
- Returns
- -------
- out : tvm.te.Tensor
- 3-D tensor with shape (batch_size, num_anchors, 6)
- """
- inter_out = multibox_transform_loc(cls_prob, loc_pred, anchor, clip, threshold, variances)
- out = non_max_suppression(
- inter_out[0],
- inter_out[1],
- inter_out[1],
- max_output_size=-1,
- iou_threshold=nms_threshold,
- force_suppress=force_suppress,
- top_k=nms_topk,
- return_indices=False,
- )
- return out
diff --git a/python/tvm/utils/__init__.py b/python/tvm/utils/__init__.py
deleted file mode 100644
index 33abc352b0..0000000000
--- a/python/tvm/utils/__init__.py
+++ /dev/null
@@ -1,19 +0,0 @@
-# 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.
-"""Utilities operating at a graph/model or other "high" level"""
-
-from .roofline import roofline_analysis
diff --git a/python/tvm/utils/roofline/__init__.py b/python/tvm/utils/roofline/__init__.py
deleted file mode 100644
index 45cc880c5b..0000000000
--- a/python/tvm/utils/roofline/__init__.py
+++ /dev/null
@@ -1,279 +0,0 @@
-# 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.
-"""Utilities for computing an approximate roofline model"""
-from typing import Dict, Optional, Union
-
-import numpy as np
-
-from ... import IRModule, auto_scheduler, build, get_global_func, nd, relay, tir, topi, transform
-from ...contrib import utils
-from ...ir.expr import GlobalVar
-from ...ir.instrument import pass_instrument
-from ...rpc.base import RPC_SESS_MASK
-from ...rpc.client import RPCSession
-from ...runtime import Device, num_threads, profiler_vm, profiling
-from ...script import tir as T
-from ...target import Target
-from . import cuda, registry, x86
-
-
-def _create_args(mod: IRModule, dev: Device, func_name: str = "main", remote=None):
- if dev.device_type >= RPC_SESS_MASK:
- random_fill = remote.get_function("tvm.contrib.random.random_fill")
- else:
- random_fill = get_global_func("tvm.contrib.random.random_fill")
- assert random_fill, "Please make sure USE_RANDOM is ON in config.cmake"
- args = []
- for arg in mod[func_name].params:
- ary = nd.empty(
- [x.value for x in arg.type_annotation.shape],
- arg.type_annotation.dtype,
- device=dev,
- )
- random_fill(ary)
- args.append(ary)
- return args
-
-
-@pass_instrument
-class SaveLoweredTIR:
- """Save TIR functions for analysis.
-
- We need the TIR function in a form that can be handled by
- `auto_scheduler.feature.named_features_from_primfunc`, but which
- is the closest to the final lowered form as possible. Right now this
- means right before tir.SplitHostDevice.
-
- """
-
- def __init__(self, before_pass: str = "tir.SplitHostDevice"):
- """
- Parameters
- ----------
- before_pass: str
- Pass before which the TIR is saved.
- """
- self.functions = {}
- self.before_pass = before_pass
-
- def run_before_pass(self, mod, info):
- if info.name == self.before_pass:
- for v, func in mod.functions.items():
- if isinstance(func, tir.PrimFunc):
- self.functions[v] = func
-
-
-def roofline_from_existing(
- report: profiling.Report,
- tir_functions: Dict[GlobalVar, tir.PrimFunc],
- target: Target,
- dev: Device,
- remote: Optional[RPCSession] = None,
-) -> profiling.Report:
- """Add roofline and other estimated statistics to an existing profiling report.
-
- :py:func:`roofline_analysis` should always be used instead of this function
- unless you need a custom compilation pipeline.
-
- Calculating roofline statistics requires features extracted the TIR
- functions in addition to per-operator runtime information (`report`) of the
- same TIR features. The features and TIR functions are not included with the
- compiled library used to generate the per-operator runtime. It is essential
- that the per-operator information comes from the exact same compilation
- pipeline as the TIR functions.
-
-
- Example
- -------
-
- ..code: : python
-
- import tvm
- import tvm.relay
-
- mod, params = tvm.relay.testing.mlp.get_workload()
-
- # it is recommended to use SaveLoweredTIR to get out the tir primfuncs
- save_tir = tvm.utils.roofline.SaveLoweredTIR()
- with tvm.transform.PassContext(opt_level=3, pass_instrument=[save_tir]):
- lib = relay.vm.compile(mod, params=params, target=target)
-
- vmexec = profiler_vm.VirtualMachineProfiler(lib, dev)
- report = vmexec.profile(*inputs)
-
- roofline_report = roofline_from_existing(report, save_tir.functions, target, dev)
-
-
- Parameters
- ----------
- report : Report
- Existing profiling report from :py:method:`VirtualMachineProfiler.profile`.
- tir_functions : Dict[GlobalVar, PrimFunc]
- TIR primfuncs from the module run to generate `report`. It is nessesary
- that these functions come before the `tir.MakePackedAPI` pass and are
- compatible with auto_scheduler featurization.
- :py:class:`SaveLoweredTIR` is the recommended way to collect these
- functions.
- target : Target
- TVM target that `report` was generated with.
- dev : Device
- Device that `report` was generated with.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
-
- Returns
- -------
- profiling.Report
- New profiling report that includes all information from `report`
- along with additional roofline metrics. See
- :py:func:`roofline_analysis` for more information on which metrics
- are included.
- """
-
- all_features = {
- prim.attrs["hash"]: (name, prim, auto_scheduler.feature.named_features_from_primfunc(prim))
- for name, prim in tir_functions.items()
- if isinstance(prim, tir.PrimFunc) and "hash" in prim.attrs.keys()
- }
-
- new_configuration = dict(report.configuration.items())
- new_calls = []
- for call in report.calls:
- if "Hash" in call.keys() and call["Hash"] in all_features:
- _, prim, features = all_features[call["Hash"]]
- if features is None:
- continue
-
- with target:
- flops, peak_flops, flops_name = registry.estimate_peak_flops(
- prim, features, target, dev, remote
- )
- loaded_bytes, peak_bandwidth, bandwidth_name = registry.estimate_peak_bandwidth(
- prim, features, target, dev, remote
- )
- new_configuration[f"Estimated Peak FLOP/s ({flops_name})"] = profiling.Ratio(peak_flops)
- new_configuration[
- f"Estimated Peak Bandwidth ({bandwidth_name}, byte/second)"
- ] = profiling.Ratio(peak_bandwidth)
- ridge_point = peak_flops / peak_bandwidth
-
- runtime = call["Duration (us)"].microseconds * 1e-6
- arith_inten = flops / loaded_bytes
- call = dict(call)
- call["Loaded Bytes"] = profiling.Count(int(loaded_bytes))
- call["Estimated FLOPs"] = profiling.Count(int(flops))
- call["Arithmetic Intensity"] = profiling.Ratio(arith_inten)
- call["FLOP/s"] = profiling.Ratio(flops / runtime)
- call["Bandwidth"] = profiling.Ratio(loaded_bytes / runtime)
- compute_bound = arith_inten > ridge_point
- call["Bound"] = "compute" if compute_bound else "memory"
- per_mem_bound = (loaded_bytes / runtime) / peak_bandwidth * 100
- per_compute_bound = (flops / runtime) / peak_flops * 100.0
- # We use ratio here because the percentages should be averaged instead of summed.
- call["Percent of Theoretical Optimal"] = profiling.Ratio(
- per_compute_bound if compute_bound else per_mem_bound
- )
- new_calls.append(call)
- else:
- new_calls.append(call)
- return profiling.Report(new_calls, report.device_metrics, new_configuration)
-
-
-def roofline_analysis(
- mod: IRModule,
- params: Dict[str, nd.NDArray],
- target: Union[str, Target],
- dev: Device,
- remote: Optional[RPCSession] = None,
-) -> profiling.Report:
- """
- Create a profiling report that contains roofline and other estimated
- statistics from running a module on the VM.
-
- The roofline model measures how close a operator gets to best possible
- memory bandwidth or FLOP/s depending on whether it is memory or compute
- bound. This computation uses the runtime of the operator along with two
- numbers extracted from the TIR code: bytes of memory touched and number of
- floating point operations.
-
- These statistics are calculated by analyzing the lowered TIR of each
- operator, so they are estimates of the true values. The statistics are:
- - Bound: Is the operator memory or compute bound. This is computed by
- assuming that the operator could perfectly cache all loads -- each byte
- of memory is only loaded once.
- - Percent of Theoretical Optimal: What percent of theoretical optimal for
- the bound. i.e. percent of peak memory bandwidth if memory bound,
- percent of peak FLOP/s if compute bound.
- - Loaded Bytes: estimation of the number of bytes loaded from main memory.
- - Estimated Flops: estimated number of floating point operations.
- - Arithmetic Intensity: ratio of FLOPs per byte of data.
- - FLOP/s: floating point operations per second.
- - Bandwidth: Number of bytes loaded per second.
-
- Parameters
- ----------
- mod : IRModule
- Uncompiled input module
-
- params : Dict[str, nd.NDArray]
-
- target : Union[str, Target]
- Target to run on.
-
- dev : Device
- Device to run on.
-
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
-
- Returns
- -------
-
- report : profiling.Report
- Profiling report which includes the estimated statistics.
- """
- if isinstance(target, str):
- target = Target(target)
-
- save_tir = SaveLoweredTIR()
- # copy existing context but add our instrument
- pass_ctx = transform.PassContext.current()
- with transform.PassContext(
- opt_level=pass_ctx.opt_level,
- required_pass=pass_ctx.required_pass,
- disabled_pass=pass_ctx.disabled_pass,
- instruments=list(pass_ctx.instruments) + [save_tir],
- config=pass_ctx.config,
- ):
- lib = relay.vm.compile(mod, params=params, target=target)
- # upload to remote if running over rpc
- if dev.device_type >= RPC_SESS_MASK:
- if remote is None:
- raise RuntimeError("A RPCSession must be provided when using a remote device.")
- temp = utils.tempdir()
- path = temp.relpath("roofline_lib.tar")
- lib.mod.export_library(path)
- remote.upload(path)
- lib = remote.load_module("roofline_lib.tar")
- vmexec = profiler_vm.VirtualMachineProfiler(lib, dev)
-
- args = _create_args(mod, dev, remote=remote)
- report = vmexec.profile(*args)
-
- return roofline_from_existing(report, save_tir.functions, target, dev, remote=remote)
diff --git a/python/tvm/utils/roofline/cuda.py b/python/tvm/utils/roofline/cuda.py
deleted file mode 100644
index b83a902b7f..0000000000
--- a/python/tvm/utils/roofline/cuda.py
+++ /dev/null
@@ -1,407 +0,0 @@
-# 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.
-"""Estimation of peak flops and memory bandwidth for cuda devices"""
-import functools
-import re
-from typing import Dict, Optional, Tuple
-
-import numpy as np
-
-from ... import build, nd, transform
-from ...contrib import nvcc, utils
-from ...rpc.base import RPC_SESS_MASK
-from ...rpc.client import RPCSession
-from ...runtime import Device
-from ...script import tir as T
-from ...target import Target
-from ...tir import PrimFunc
-from . import registry
-
-
-@functools.lru_cache(maxsize=None)
-def estimate_peak_flops_tensorcore(
- target: Target,
- dev: Device,
- remote: Optional[RPCSession],
- mat_dtype: str = "float16",
- acc_dtype: str = "float32",
-) -> Tuple[float, float, str]:
- """Estimate the peak FLOP/s of a cuda device with tensorcores.
-
- This estimate should only be used to compare with operators that can use
- dense tensorcore mma instructions.
-
- References
- ----------
- Wei Sun, Ang Li, Tong Geng, Sander Stuijk, Henk Corporaal: "Dissecting
- Tensor Cores via Microbenchmarks: Latency, Throughput and Numerical
- Behaviors", 2022; http://arxiv.org/abs/2206.02874
- https://www.nvidia.com/content/PDF/nvidia-ampere-ga-102-gpu-architecture-whitepaper-v2.1.pdf
-
- Parameters
- ----------
- target : Target
- Target to run on. This should be as specific to the actual hardware as
- possible.
- dev : Device
- Device to run on.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
- mat_dtype : str
- Dtype of matrices passed to mma instructions.
- acc_dtype : str
- Dtype of accumulator to use with mma instructions. Should be compatible
- with `mat_dtype`.
-
- Returns
- -------
- peak_flops : float
- Approximate sustained FLOP/s of this target/device combo assuming
- mma instructions. Addition and multiplications are each counted as
- separate FLOPs.
- """
-
- @T.prim_func
- def peak_flops_tensorcore_tir(
- inp: T.Buffer((16, 16), mat_dtype),
- out: T.Buffer((16, 16), acc_dtype),
- n: T.int32,
- sms: T.int32,
- ):
- # pylint: disable=invalid-name, missing-function-docstring
- A = T.alloc_buffer((16, 16), dtype=mat_dtype, scope="wmma.matrix_a")
- B = T.alloc_buffer((16, 16), dtype=mat_dtype, scope="wmma.matrix_b")
- C = T.alloc_buffer((16, 16), dtype=acc_dtype, scope="wmma.accumulator")
- for _ in T.thread_binding(sms, thread="blockIdx.x"):
- for _ in T.thread_binding(
- 8, thread="threadIdx.y"
- ): # need 8 warps to get enough in-SM parallelism
- for _ in T.thread_binding(32, thread="threadIdx.x"):
- T.evaluate(
- T.tvm_load_matrix_sync(
- A.data,
- 16,
- 16,
- 16,
- 0,
- T.tvm_access_ptr(
- T.type_annotation(dtype=mat_dtype),
- inp.data,
- 0,
- 16,
- 1,
- dtype="handle",
- ),
- 16,
- "row_major",
- dtype="handle",
- )
- )
- T.evaluate(T.tvm_fill_fragment(B.data, 16, 16, 16, 0, 0, dtype="handle"))
- T.evaluate(T.tvm_fill_fragment(C.data, 16, 16, 16, 0, 0, dtype="handle"))
- for _ in range(n):
- T.evaluate(
- T.tvm_mma_sync(
- C.data, 0, A.data, 0, B.data, 0, C.data, 0, dtype="handle"
- )
- )
- T.evaluate(
- T.tvm_store_matrix_sync(
- C.data,
- 16,
- 16,
- 16,
- 0,
- T.tvm_access_ptr(
- T.type_annotation(dtype=acc_dtype),
- out.data,
- 0,
- 16,
- 2,
- dtype="handle",
- ),
- 16,
- "row_major",
- dtype="handle",
- )
- )
-
- n = 100000
- sms = dev.multi_processor_count
- specialized = peak_flops_tensorcore_tir.specialize(
- {peak_flops_tensorcore_tir.params[2]: n, peak_flops_tensorcore_tir.params[3]: sms}
- )
- with transform.PassContext(opt_level=3):
- f = build(specialized, target=target)
-
- # upload to remote if running over rpc
- if dev.device_type >= RPC_SESS_MASK:
- if remote is None:
- raise RuntimeError("A RPCSession must be provided when using a remote device.")
- temp = utils.tempdir()
- path = temp.relpath("peak_mma_flops.tar")
- f.export_library(path)
- remote.upload(path)
- f = remote.load_module("peak_mma_flops.tar")
-
- x = nd.empty((16, 16), dtype=mat_dtype, device=dev)
- y = nd.empty((16, 16), dtype=acc_dtype, device=dev)
- times = f.time_evaluator(f.entry_name, dev, repeat=10, number=1)(x, y)
- # each mma operation computes 16 x 16 x 16 FLOPs
- return n * 16 * 16 * 16 * 2 * sms * 8 / times.min
-
-
-@functools.lru_cache(maxsize=None)
-def estimate_peak_flops_fma(
- target: Target,
- dev: Device,
- remote: Optional[RPCSession],
- dtype: str,
-) -> Tuple[float, float, str]:
- """Estimate the peak FLOP/s of a cuda device with fma operations (not using tensor cores).
-
- References
- ----------
- https://www.nvidia.com/content/PDF/nvidia-ampere-ga-102-gpu-architecture-whitepaper-v2.1.pdf
-
- Parameters
- ----------
- target : Target
- Target to run on. This should be as specific to the actual hardware as
- possible.
- dev : Device
- Device to run on.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
- dtype : str
- Dtype of fma operation
-
- Returns
- -------
- peak_flops : float
- Approximate sustained FLOP/s of this target/device combo assuming
- fma instructions. Addition and multiplications are each counted as
- separate FLOPs.
- """
-
- vec_width = 32
- warps = 16 # need 16 warps to get enough in-SM parallelism
- sms = dev.multi_processor_count
- n = 100000
-
- @T.prim_func
- def peak_flops_fma_tir(
- A: T.Buffer((sms, warps, vec_width), dtype),
- B: T.Buffer((sms, warps, vec_width), dtype),
- ):
- # pylint: disable=invalid-name, missing-function-docstring
- shared = T.alloc_buffer((sms, warps, vec_width), dtype=dtype, scope="shared")
- for sm in T.thread_binding(sms, thread="blockIdx.x"):
- for warp in T.thread_binding(warps, thread="threadIdx.y"):
- for t in T.thread_binding(vec_width, thread="threadIdx.x"):
- shared[sm, warp, t] = A[sm, warp, t]
- for _ in range(n):
- shared[sm, warp, t] = (
- shared[sm, warp, t] * shared[sm, warp, t] + shared[sm, warp, t]
- )
- B[sm, warp, t] = shared[sm, warp, t]
-
- with transform.PassContext(opt_level=3):
- f = build(peak_flops_fma_tir, target=target)
-
- # upload to remote if running over rpc
- if dev.device_type >= RPC_SESS_MASK:
- if remote is None:
- raise RuntimeError("A RPCSession must be provided when using a remote device.")
- temp = utils.tempdir()
- path = temp.relpath("peak_fma_flops.tar")
- f.export_library(path)
- remote.upload(path)
- f = remote.load_module("peak_fma_flops.tar")
-
- x = nd.empty((sms, warps, vec_width), dtype=dtype, device=dev)
- y = nd.empty((sms, warps, vec_width), dtype=dtype, device=dev)
- times = f.time_evaluator(f.entry_name, dev, repeat=10, number=1)(x, y)
- return n * warps * sms * vec_width * 2 / times.min
-
-
-@registry.estimate_peak_flops.register("cuda")
-def estimate_peak_flops(
- func: PrimFunc, # pylint: disable=unused-argument
- features: Dict[str, np.ndarray],
- target: Target,
- dev: Device,
- remote: Optional[RPCSession],
-) -> Tuple[float, float, str]:
- """Estimate the peak FLOP/s of a cuda device.
-
- Parameters
- ----------
- func : PrimFunc
- Function to estimate peak flops for. Used to check if a specific kind
- intrinsic or dtype could be used with this function.
- features : Dict[str, np.ndarry]
- Features extracted from `func`. Used to check if a specific kind
- intrinsic or dtype could be used with this function.
- target : Target
- Target to run on. This should be as specific to the actual hardware as
- possible.
- dev : Device
- Device to run on.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
-
- Returns
- -------
- flops : float
- Estimated number of flops used by `func`.
- peak_flops : float
- Approximate sustained FLOP/s of this target/device combo. Addition and
- multiplications are each counted as separate FLOPs.
- name : str
- Dtype/intrinsic used by `func` to achieve peak flops.
- """
- has_tensorcore = nvcc.have_tensorcore(dev.compute_version)
- # assume that the first argument dtype is the same as all the others
- dtype = list(func.buffer_map.values())[0].dtype
- if dtype == "float16" and has_tensorcore:
- peak_flops = estimate_peak_flops_tensorcore(target, dev, remote)
- name = "float16 tensorcore"
- else:
- peak_flops = estimate_peak_flops_fma(target, dev, remote, dtype)
- name = f"{dtype} fma"
- flops = np.sum(
- features["float_addsub"]
- + features["float_mul"]
- + features["float_mad"] * 2
- + features["float_divmod"]
- )
- return flops, peak_flops, name
-
-
-@T.prim_func
-def peak_bandwidth_tir(a: T.handle, b: T.handle, blocks: T.int32, warp_size: T.int32) -> None:
- # pylint: disable=invalid-name, missing-function-docstring
- N = T.int32()
- A = T.match_buffer(a, [blocks, N, 4, warp_size], "float32")
- B = T.match_buffer(b, [blocks, 4, warp_size], "float32")
- for i in T.thread_binding(blocks, "blockIdx.x"):
- for k in T.serial(N):
- for l in T.unroll(4):
- # vectorized load is necessary to hit peak bandwidth
- for j in T.thread_binding(warp_size, "threadIdx.x"):
- # += is necessary to introduce a data dependency for all
- # elements of A, preventing the backend from removing the
- # `k` loop and setting `k` to the loop extent.
- B[i, l, j] += A[i, k, l, j]
-
-
-@functools.lru_cache(maxsize=None)
-def estimate_peak_bandwidth_global_mem(
- target: Target,
- dev: Device,
- remote: Optional[RPCSession] = None,
-) -> Tuple[float, float, str]:
- """Estimate peak bandwidth of global memory. See estimate_peak_bandwidth"""
- warp_size = dev.warp_size
- # These sizes seem large enough to give the card time to hit a fixpoint on memory bandwidth
- blocks = 1024
- size = 1024
-
- specialized = peak_bandwidth_tir.specialize(
- {peak_bandwidth_tir.params[2]: blocks, peak_bandwidth_tir.params[3]: warp_size}
- )
- with transform.PassContext(opt_level=3):
- f = build(specialized, target=target)
-
- # upload to remote if running over rpc
- if dev.device_type >= RPC_SESS_MASK:
- if remote is None:
- raise RuntimeError("A RPCSession must be provided when using a remote device.")
- temp = utils.tempdir()
- path = temp.relpath("peak_bandwidth.tar")
- f.export_library(path)
- remote.upload(path)
- f = remote.load_module("peak_bandwidth.tar")
-
- a = nd.empty((blocks, size, 4, warp_size), dtype="float32", device=dev)
- b = nd.empty((blocks, 4, warp_size), dtype="float32", device=dev)
- times = f.time_evaluator(f.entry_name, dev, repeat=10, number=1)(a, b)
- return a.numpy().size * 4 / times.min # 4 bytes per float32
-
-
-@registry.estimate_peak_bandwidth.register("cuda")
-def estimate_peak_bandwidth(
- func: PrimFunc, # pylint: disable=unused-argument
- features: Dict[str, np.ndarray],
- target: Target,
- dev: Device,
- remote: Optional[RPCSession] = None,
-) -> Tuple[float, float, str]:
- """Estimate peak memory bandwidth of a target/device combo.
-
- Peak bandwidth is estimated by running a small experiment on the underlying
- hardware. The peak bandwidth measurement assumes that vector instructions
- are being used to load the data.
-
- Parameters
- ----------
- func : PrimFunc
- Function to estimate peak bandwidth for. Used to check if a specific
- kind of memory could be used with this function.
- features : Dict[str, np.ndarry]
- Features extracted from `func`. Used to check if a specific kind of
- memory could be used with this function.
- target : Target
- Target to use for measurement. This target should be as specific to the
- underlying hardware as possible.
- dev : Device
- Device to measure peak bandwidth on.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
-
- Returns
- -------
- loaded_bytes : float
- Estimated bytes loaded by `func`.
- peak_bandwidth : float
- Peak memory bandwidth in bytes/seconds.
- name : str
- Name of the memory being used.
- """
- # autoscheduler features do not take into account that 1.
- # global and shared memory have very different performance
- # characteristics -- both are included in the same bytes
- # touched count 2. multiple threads accessing the same byte
- # of memory does not use the same amount of bandwidth as
- # multiple threads accessing different bytes of memory. We
- # use unique bytes accessed here to avoid these two issues,
- # but this does bias results towards being more compute
- # bound.
- loaded_bytes = sum(
- [
- np.sum(x)
- for (k, x) in features.items()
- if re.match(r"^B[0-9]+\.unique_bytes$", k) is not None
- ]
- )
- peak_bandwidth = estimate_peak_bandwidth_global_mem(target, dev, remote)
- return loaded_bytes, peak_bandwidth, "global"
diff --git a/python/tvm/utils/roofline/registry.py b/python/tvm/utils/roofline/registry.py
deleted file mode 100644
index 9358529b38..0000000000
--- a/python/tvm/utils/roofline/registry.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# 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.
-"""Definition of generic functions for estimating peak flops and bandwidth"""
-from typing import Dict, Optional, Tuple
-
-import numpy as np
-
-from ...rpc.client import RPCSession
-from ...runtime import Device
-from ...target import Target, generic_func
-from ...tir import PrimFunc
-
-
-@generic_func
-def estimate_peak_bandwidth(
- func: PrimFunc,
- features: Dict[str, np.ndarray],
- target: Target,
- dev: Device,
- remote: Optional[RPCSession] = None,
-) -> Tuple[float, float, str]:
- """Estimate peak memory bandwidth of a target/device combo.
-
- Peak bandwidth is estimated by running a small experiment on the underlying
- hardware. The peak bandwidth measurement assumes that vector instructions
- are being used to load the data.
-
- Parameters
- ----------
- func : PrimFunc
- Function to estimate peak bandwidth for. Used to check if a specific
- kind of memory could be used with this function.
- features : Dict[str, np.ndarry]
- Features extracted from `func`. Used to check if a specific kind of
- memory could be used with this function.
- target : Target
- Target to use for measurement. This target should be as specific to the
- underlying hardware as possible.
- dev : Device
- Device to measure peak bandwidth on.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
-
- Returns
- -------
- loaded_bytes : float
- Estimated bytes loaded by `func`.
- peak_bandwidth : float
- Peak memory bandwidth in bytes/seconds.
- name : str
- Name of the memory being used.
- """
- raise NotImplementedError()
-
-
-@generic_func
-def estimate_peak_flops(
- func: PrimFunc,
- features: Dict[str, np.ndarray],
- target: Target,
- dev: Device,
- remote: Optional[RPCSession],
-) -> Tuple[float, float, str]:
- """
- Estimate the maximum number of FLOP/s this target/device combo is capable
- of reaching by running a test program. This is a generic function that
- should be overridden for each target.
-
- Parameters
- ----------
- func : PrimFunc
- Function to estimate peak flops for. Used to check if a specific kind
- intrinsic or dtype could be used with this function.
- features : Dict[str, np.ndarry]
- Features extracted from `func`. Used to check if a specific kind
- intrinsic or dtype could be used with this function.
- target : Target
- Target to run on. This should be as specific to the actual hardware as
- possible to make sure that LLVM generates the best vector code.
- dev : Device
- Device to run on.
- remote : Optional[RPCSession]
- Remote session used to upload artifacts for runtime evaluation. Must be
- the same session used to create `dev`.
-
- Returns
- -------
- flops : float
- Estimated number of flops used by `func`.
- peak_flops : float
- Approximate sustained FLOP/s of this target/device combo assuming
- vectorized FMA instructions. Each FMA operation counts as two FLOPs.
- name : str
- Dtype/intrinsic used by `func` to achieve peak flops.
- """
- raise NotImplementedError()
... 15290 lines suppressed ...