[TIR] Handle subroutine calls in MakePackedAPI #14913

Lunderberg · 2023-05-22T13:48:29Z

Previously, the MakePackedAPI lowering pass treated all functions as externally-visible functions, which require the "global_symbol" attribute and are updated to use the PackedFunc API. This commit updates MakePackedAPI to be aware of subroutines, distinguishing between externally-visible subroutines, which have the "global_symbol" attribute, and internal subroutines, which do not have the "global_symbol" attribute.

Externally-visible subroutines may be called directly by the user of a runtime.Module, and must use the packed func interface. Any internal callers of an externally-visible subroutine must be updated to use the builtin::tvm_call_cpacked().

Internal subroutines may not be called directly by an external user, and neither the subroutine's signature nor any internal callers should be updated to the packed func interface.

Previously, `MakePackedAPI` required all functions to have the `kGlobalSymbol` attribute. This commit updates the behavior such that `MakePackedAPI` only modifies PrimFuncs that have the `kGlobalSymbol` attribute, and passes through any other PrimFunc unmodified.

When a function is updated to use the `PackedFunc` API, any calls made to that function from elsewhere in the `IRModule` should be updated as well.

tvm-bot · 2023-05-22T13:48:33Z

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

Lunderberg · 2023-05-22T13:49:06Z

This PR is a subset of the functionality in #14862. As that PR has grown to a difficult-to-review size, I'm carving off independent chunks of it into separate PRs, after which #14862 can be rebased to a more manageable size.

Prior to this commit, MakeUnpackedAPI required all functions to be annotated with `kGlobalSymbol` (`"global_symbol"`). This commit updates the transformation to apply only to functions that have the `kGlobalSymbol` attribute, and to update any internal callsites of the modified functions. This is analogous to the changes made in apache#14913, which updates `MakePackedAPI`.

Prior to this commit, MakeUnpackedAPI required all functions to be annotated with `kGlobalSymbol` (`"global_symbol"`). This commit updates the transformation to apply only to functions that have the `kGlobalSymbol` attribute, and to update any internal callsites of the modified functions. This is analogous to the changes made in #14913, which updates `MakePackedAPI`.

PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict.

…#14950) PRs #14913 and #14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict.

) * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs #14913 and #14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [Target] Added WithoutHost method * [TIR] Preserve existing kTarget function attribute in BindTarget Previously, if a function already has a `tvm::attr::kTarget` attribute, it will be overwritten by the `tir.BindTarget` transform. This commit updates the behavior such that `tir.BindTarget` adds annotations to functions that are missing a target annotation, but preserves any existing target annotations. This is part of a series of commits to simplify the handling of multi-target builds.

) * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs #14913 and #14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [Target] Added utility method TargetNode::HasKey() This utility method makes it easier to determine if a target contains a specific key. * [TIR] Added utility method tvm::tir::IsHostFunc(const PrimFunc&) For modules that contain both host and device functions, this utility function checks whether a given PrimFunc is a host function, based on the target annotation. * [TIR] Restrict InstallDebugSpans to host functions Previously, the `tir.InstallDebugSpans` pass required the module to contain only a single PrimFunc. This commit relaxes the requirement, to require a single host-side PrimFunc, and to ignore any other device-side functions.

) * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs #14913 and #14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [TVMScript] Prevent bool to int conversion in T.Assert condition Previously, while literal `True` and `False` values were converted to `tvm::Bool` instances, constant-foldable expressions (e.g. `0 == 1`) would be evaluated to `True`, but were then passed directly to the FFI. Because the FFI uses the same representation for integer and boolean values, the conversion to `PrimExpr` resulted in a `tvm::Integer` instead of `tvm::Bool`. This commit converts the argument of `T.Assert` to a `tvm::Bool` before calling the FFI, avoiding the ambiguity. In addition, the `AssertStmt` constructor now validates the datatype of the condition, to prevent it from re-occurring. This was first caught in the unit test `test_debug_info.py::test_llvm_ir_debug_info`, which failed on some versions of LLVM due to the use of `i32` as the condition of an assert. * Updated TVMScript printer unit test with boolean condition

* [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs #14913 and #14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [Codegen][LLVM] Allow void return type from PackedFunc Previously, calling a packed func that returns void would result in an error being raised from `tir::APIType`, as there is no runtime representation of a void type. This commit updates `CodeGenCPU::MakeCallPackedLowered` to only read the return value and type fo a `PackedFunc` when the TIR return type is non-void.

This PR refactors SplitHostDevice into three separate transformations. Previously, SplitHostDevice would replace device regions with a builtin::tvm_call_packed() node to replace the extracted region. After this PR, this process is performed in three separate steps. AnnotateDeviceRegion: Annotate the regions that should be executed on another target. SplitHostDevice: Extract the annotated region into an independent PrimFunc, with a GlobalVar to represent the call from into the new subroutine. LowerDeviceKernelLaunch: For any subroutine call where the caller and callee are on different devices, replace with a device kernel launch. * PR#14915 [TVMScript] Allow T.target("device", host="host") in TVMScript Prior to this commit, the `TargetNode::host` could be specified in TVMScript as part of the config dictionary, under the key `"host"`. However, this required all other device parameters to be explicitly specified, rather than using any of the short-hand string representations. This commit forwards the `host` argument from TVMScript's `T.target` method to `tvm.target.Target`, allowing both the device and host to be specified using the shorthand string representation. ```python @T.prim_func def before_this_commit(): T.func_attr( { "target": T.target( { "arch": "sm_86", "host": {"keys": ["cpu"], "kind": "llvm", "tag": ""}, "keys": ["cuda", "gpu"], "kind": "cuda", "max_num_threads": 1024, "tag": "", "thread_warp_size": 32, } ) } ) T.evaluate(0) @T.prim_func def after_this_commit(): T.func_attr({"target": T.target("cuda", host="llvm")}) T.evaluate(0) ``` * [Target] Added WithoutHost method * [TIR] SplitHostDevice, handle missing kGlobalSymbol Previously, the symbol name of the extracted compute kernel was defined based on the `kGlobalSymbol` attribute, which was required to be present. This commit updates `SplitHostDevice` to generate the symbol name using `kGlobalSymbol` if present, and to fall back to the name of the `tvm::GlobalVar` for internal functions. * [TIR] Refactor SplitHostDevice into three separate passes First pass, `AnnotateDeviceRegions`. This pass decides which portions of a PrimFunc should be run on the device, and annotates them with `kTarget` attribute, indicating which target should be used for later lowering steps. Second pass, `SplitHostDevice`. This pass extracts the annotated region into an independent PrimFunc. The `kTarget` attribute of the extracted kernel is defined by the `kTarget` annotation inserted by `AnnotateDeviceRegions`. The host function is marked by the `tvm::tir::attr::kIsHostFunc` attribute, allowing it to be recognized by later host-only lowering passes. Third pass, `LowerDeviceKernelLaunch`. This pass identifies subroutine calls that call into device kernels, and rewrites them into `T.tvm_call_packed`. * Add unit tests specifically for SplitHostDevice behavior * Added unit test specifically for AnnotateDeviceRegions * Added unit tests for LowerDeviceKernelLaunch * Minor cleanup, moved all kernel launch collection into one spot Previously, the SplitHostDevice pass added the `tir::attr::kKernelLaunchParams` attribute, and the LowerDeviceKernelLaunch pass filled in the values for it. This cleanup makes the kernel launch params be the sole responsibility of LowerDeviceKernelLaunch. * Updated unit tests for LowerWarpMemory * Updated unit tests for ThreadSync * Updated unit test for inject ptx async copy * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs #14913 and #14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * Maintain "tir.is_global_func" attr in device-side entry point * SplitHostDevice, update the host-side target to be the host * [TIR] Update LowerDeviceKernelLaunch to avoid kIsHostFunc Update to use the `tvm::tir::IsHostFunc` utility function, rather than the `kIsHostFunc` attribute. Per discussion on #14020, the `kIsHostFunct` attribute should only be used in `BindTarget`, and should not be re-introduced in `SplitHostDevice`. * Remove is_host_func from SplitHostDevice tests

The functionality tested in this commit was added across several recent PRs, each of which tested their features in isolation. This PR adds unit tests to validate the end-to-end behavior of TIR subroutine calls. PRs building up to this point: - TVMScript - apache#14889 - apache#14915 - apache#14919 - apache#14941 - Functionality improvements of existing TIR passes - apache#14913 - apache#14914 - apache#14918 - apache#14951 - Changes to the TIR lowering flow - apache#14942 - apache#14985 - Codegen updates - apache#14958 - apache#14901 - Compatibility updates/fixes - apache#14892 - apache#14950 - apache#14943 - apache#14944 - apache#14945 - apache#14952 - apache#14982 - apache#14949

* [TIR] MakePackedAPI, handle missing kGlobalSymbol Previously, `MakePackedAPI` required all functions to have the `kGlobalSymbol` attribute. This commit updates the behavior such that `MakePackedAPI` only modifies PrimFuncs that have the `kGlobalSymbol` attribute, and passes through any other PrimFunc unmodified. * [TIR] Update calls to externally-exposed subroutines in MakePackedAPI When a function is updated to use the `PackedFunc` API, any calls made to that function from elsewhere in the `IRModule` should be updated as well. * Bugfix, don't update the callsite unless the callee is also updated

Prior to this commit, MakeUnpackedAPI required all functions to be annotated with `kGlobalSymbol` (`"global_symbol"`). This commit updates the transformation to apply only to functions that have the `kGlobalSymbol` attribute, and to update any internal callsites of the modified functions. This is analogous to the changes made in apache#14913, which updates `MakePackedAPI`.

…apache#14950) PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict.

…che#14942) * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [Target] Added WithoutHost method * [TIR] Preserve existing kTarget function attribute in BindTarget Previously, if a function already has a `tvm::attr::kTarget` attribute, it will be overwritten by the `tir.BindTarget` transform. This commit updates the behavior such that `tir.BindTarget` adds annotations to functions that are missing a target annotation, but preserves any existing target annotations. This is part of a series of commits to simplify the handling of multi-target builds.

…che#14943) * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [Target] Added utility method TargetNode::HasKey() This utility method makes it easier to determine if a target contains a specific key. * [TIR] Added utility method tvm::tir::IsHostFunc(const PrimFunc&) For modules that contain both host and device functions, this utility function checks whether a given PrimFunc is a host function, based on the target annotation. * [TIR] Restrict InstallDebugSpans to host functions Previously, the `tir.InstallDebugSpans` pass required the module to contain only a single PrimFunc. This commit relaxes the requirement, to require a single host-side PrimFunc, and to ignore any other device-side functions.

…che#14941) * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [TVMScript] Prevent bool to int conversion in T.Assert condition Previously, while literal `True` and `False` values were converted to `tvm::Bool` instances, constant-foldable expressions (e.g. `0 == 1`) would be evaluated to `True`, but were then passed directly to the FFI. Because the FFI uses the same representation for integer and boolean values, the conversion to `PrimExpr` resulted in a `tvm::Integer` instead of `tvm::Bool`. This commit converts the argument of `T.Assert` to a `tvm::Bool` before calling the FFI, avoiding the ambiguity. In addition, the `AssertStmt` constructor now validates the datatype of the condition, to prevent it from re-occurring. This was first caught in the unit test `test_debug_info.py::test_llvm_ir_debug_info`, which failed on some versions of LLVM due to the use of `i32` as the condition of an assert. * Updated TVMScript printer unit test with boolean condition

* [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * [Codegen][LLVM] Allow void return type from PackedFunc Previously, calling a packed func that returns void would result in an error being raised from `tir::APIType`, as there is no runtime representation of a void type. This commit updates `CodeGenCPU::MakeCallPackedLowered` to only read the return value and type fo a `PackedFunc` when the TIR return type is non-void.

This PR refactors SplitHostDevice into three separate transformations. Previously, SplitHostDevice would replace device regions with a builtin::tvm_call_packed() node to replace the extracted region. After this PR, this process is performed in three separate steps. AnnotateDeviceRegion: Annotate the regions that should be executed on another target. SplitHostDevice: Extract the annotated region into an independent PrimFunc, with a GlobalVar to represent the call from into the new subroutine. LowerDeviceKernelLaunch: For any subroutine call where the caller and callee are on different devices, replace with a device kernel launch. * PR#14915 [TVMScript] Allow T.target("device", host="host") in TVMScript Prior to this commit, the `TargetNode::host` could be specified in TVMScript as part of the config dictionary, under the key `"host"`. However, this required all other device parameters to be explicitly specified, rather than using any of the short-hand string representations. This commit forwards the `host` argument from TVMScript's `T.target` method to `tvm.target.Target`, allowing both the device and host to be specified using the shorthand string representation. ```python @T.prim_func def before_this_commit(): T.func_attr( { "target": T.target( { "arch": "sm_86", "host": {"keys": ["cpu"], "kind": "llvm", "tag": ""}, "keys": ["cuda", "gpu"], "kind": "cuda", "max_num_threads": 1024, "tag": "", "thread_warp_size": 32, } ) } ) T.evaluate(0) @T.prim_func def after_this_commit(): T.func_attr({"target": T.target("cuda", host="llvm")}) T.evaluate(0) ``` * [Target] Added WithoutHost method * [TIR] SplitHostDevice, handle missing kGlobalSymbol Previously, the symbol name of the extracted compute kernel was defined based on the `kGlobalSymbol` attribute, which was required to be present. This commit updates `SplitHostDevice` to generate the symbol name using `kGlobalSymbol` if present, and to fall back to the name of the `tvm::GlobalVar` for internal functions. * [TIR] Refactor SplitHostDevice into three separate passes First pass, `AnnotateDeviceRegions`. This pass decides which portions of a PrimFunc should be run on the device, and annotates them with `kTarget` attribute, indicating which target should be used for later lowering steps. Second pass, `SplitHostDevice`. This pass extracts the annotated region into an independent PrimFunc. The `kTarget` attribute of the extracted kernel is defined by the `kTarget` annotation inserted by `AnnotateDeviceRegions`. The host function is marked by the `tvm::tir::attr::kIsHostFunc` attribute, allowing it to be recognized by later host-only lowering passes. Third pass, `LowerDeviceKernelLaunch`. This pass identifies subroutine calls that call into device kernels, and rewrites them into `T.tvm_call_packed`. * Add unit tests specifically for SplitHostDevice behavior * Added unit test specifically for AnnotateDeviceRegions * Added unit tests for LowerDeviceKernelLaunch * Minor cleanup, moved all kernel launch collection into one spot Previously, the SplitHostDevice pass added the `tir::attr::kKernelLaunchParams` attribute, and the LowerDeviceKernelLaunch pass filled in the values for it. This cleanup makes the kernel launch params be the sole responsibility of LowerDeviceKernelLaunch. * Updated unit tests for LowerWarpMemory * Updated unit tests for ThreadSync * Updated unit test for inject ptx async copy * [Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI PRs apache#14913 and apache#14914 made analogous changes to `MakePackedAPI` and `MakeUnpackedAPI` to handle subroutine calls. Both PRs introduced the same symbol, `tvm::tir::SubroutineCallRewriter`, a local utility to update internal calls to a modified function. While each PR passed CI individually, and was therefore able to merge, having both changes caused a duplicate symbol. This commit updates `MakePackedAPI` and `MakeUnpackedAPI` to place their local utilities into anonymous namespaces, avoiding the conflict. * Maintain "tir.is_global_func" attr in device-side entry point * SplitHostDevice, update the host-side target to be the host * [TIR] Update LowerDeviceKernelLaunch to avoid kIsHostFunc Update to use the `tvm::tir::IsHostFunc` utility function, rather than the `kIsHostFunc` attribute. Per discussion on apache#14020, the `kIsHostFunct` attribute should only be used in `BindTarget`, and should not be re-introduced in `SplitHostDevice`. * Remove is_host_func from SplitHostDevice tests

The functionality tested in this commit was added across several recent PRs, each of which tested their features in isolation. This PR adds unit tests to validate the end-to-end behavior of TIR subroutine calls. PRs building up to this point: - TVMScript - apache#14889 - apache#14915 - apache#14919 - apache#14941 - Functionality improvements of existing TIR passes - apache#14913 - apache#14914 - apache#14918 - apache#14951 - Changes to the TIR lowering flow - apache#14942 - apache#14985 - Codegen updates - apache#14958 - apache#14901 - Compatibility updates/fixes - apache#14892 - apache#14950 - apache#14943 - apache#14944 - apache#14945 - apache#14952 - apache#14982 - apache#14949

Lunderberg added 2 commits May 22, 2023 08:37

[TIR] Update calls to externally-exposed subroutines in MakePackedAPI

d706d35

When a function is updated to use the `PackedFunc` API, any calls made to that function from elsewhere in the `IRModule` should be updated as well.

Lunderberg mentioned this pull request May 22, 2023

[TIR] Handle subroutine calls in MakeUnpackedAPI #14914

Merged

Bugfix, don't update the callsite unless the callee is also updated

9859694

csullivan approved these changes May 24, 2023

View reviewed changes

Lunderberg merged commit 41a616f into apache:main May 25, 2023
18 checks passed

Lunderberg deleted the make_packed_api_handle_subroutines branch May 25, 2023 12:37

Lunderberg mentioned this pull request May 25, 2023

[Bugfix] Avoid symbol conflicts in MakePackedAPI/MakeUnpackedAPI #14950

Merged

ysh329 mentioned this pull request Jul 12, 2023

[Release] v0.13.0 Release Candidate Notes #15295

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TIR] Handle subroutine calls in MakePackedAPI #14913

[TIR] Handle subroutine calls in MakePackedAPI #14913

Lunderberg commented May 22, 2023

tvm-bot commented May 22, 2023

Lunderberg commented May 22, 2023

[TIR] Handle subroutine calls in MakePackedAPI #14913

[TIR] Handle subroutine calls in MakePackedAPI #14913

Conversation

Lunderberg commented May 22, 2023

tvm-bot commented May 22, 2023

Lunderberg commented May 22, 2023