Lunderberg commented on code in PR #16721:
URL: https://github.com/apache/tvm/pull/16721#discussion_r1526229596
##########
src/tir/transforms/lower_tvm_builtin.cc:
##########
@@ -38,6 +38,19 @@ namespace tir {
// These information are needed during codegen.
class BuiltinLower : public StmtExprMutator {
public:
+ static PrimFunc Build(PrimFunc func) {
Review Comment:
This was another bug I ran into. The `LowerTVMBuiltin` pass would look for
each `Allocate` node, and would replace it with a call to the PackedFunc
`TVMBackendAllocWorkspace`, using the exact device type on which the allocation
should occur. This change allowed it to search in the function attributes for
the device type, and not just in the `AttrStmt`.
Partly, this is due to the expanded role of the `"tir.is_host_func"`
attribute. Initially, this attribute was used to tag the host/device portions
after `SplitHostDevice`. However, it also has been useful for tagging a
function as being on the host, when in a portion of the lowering flow that
doesn't yet know the exact targets for the device/host
([example](https://github.com/apache/tvm/blob/main/src/relax/backend/vm/vm_shape_lower.cc#L593)).
This is useful functionality, so I expanding `LowerTVMBuiltin` to support it
makes sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
