search

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.46k stars 3.41k forks source link

[Relax][Bugfix] Annotate ComputePrimValue output as host function #17032

Closed Lunderberg closed 2 months ago

Lunderberg commented 2 months ago

The ComputePrimValue transform is used to compute the value of symbolic expressions that may appear within a Relax function. For example, to compute a boolean condition used for a relax::If node. These functions are used for small host-side computations, prior to launching a device kernel.

This commit updates ComputePrimValue to annotate the generated PrimFunc with tir::attr::kIsHostFunc. This annotation is required for correct behavior in tvm.dlight.ApplyDefaultSchedule, to avoid erroneous scheduling of this function for the GPU, and for tir::transform::BindTarget, to ensure that the function is compiled for execution on the host.

Lunderberg commented 2 months ago

This PR is identical to the now-closed PR https://github.com/apache/tvm/pull/17024, as a CI restart prevented the CI from running on the original PR.