tqchen commented on code in PR #17599:
URL: https://github.com/apache/tvm/pull/17599#discussion_r2014592454
##########
python/tvm/relax/transform/transform.py:
##########
@@ -1605,6 +1611,45 @@ def AllocateWorkspace() -> tvm.ir.transform.Pass:
return _ffi_api.AllocateWorkspace() # type: ignore
+def AnnotateCustomMemoryScope(target: Optional[Target] = None) ->
tvm.ir.transform.Pass:
+ """Allocate the memory scope information. This is Adreno specific pass to
annotate
+ The memory scope information and realize the same with RealizeVDevice pass
followed by
+ updating the Prim Function var_buffer mapping using
SpecializePrimFuncBasedOnCallSite.
+
+ Returns
+ -------
+ ret: tvm.ir.transform.Pass
+ The registered pass for allocating workspace.
+ """
+ return _ffi_api.AnnotateCustomMemoryScope(target) # type: ignore
+
+
+def SpecializePrimFuncBasedOnCallSite() -> tvm.ir.transform.Pass:
+ """This pass updates the var_buffer mapping of PrimFunctions from the
call_tir info.
+ Primarily used to update the VDevice information if any changes occured
from the caller.
+ This pass recreates the buffers and updates the map.
+
+ Returns
+ -------
+ ret: tvm.ir.transform.Pass
+ The registered pass for allocating workspace.
+ """
+ return _ffi_api.SpecializePrimFuncBasedOnCallSite() # type: ignore
+
+
+def FoldVDeviceScopeChange() -> tvm.ir.transform.Pass:
Review Comment:
move to backend.andreno namespace
##########
src/relax/transform/legalize_ops.cc:
##########
@@ -152,6 +164,68 @@ class LegalizeMutator : public ExprMutator {
return NullOpt;
}
+ Expr UpdateVDeviceOutStructInfo(Expr expr, const Call& visited_call) {
+ static const auto& infer_struct_info_map =
Op::GetAttrMap<FInferStructInfo>("FInferStructInfo");
+ static const Op& call_tir_op = Op::Get("relax.call_tir");
+ auto* op_node = visited_call->op.as<OpNode>();
+
+ // Not an OpNode
+ if (op_node == nullptr) {
+ return expr;
+ }
+ auto op = GetRef<Op>(op_node);
+
+ if (!infer_struct_info_map.count(op)) {
+ return expr;
+ }
+
+ if (!expr->IsInstance<CallNode>()) {
+ return expr;
+ }
+
+ auto call = Downcast<Call>(expr);
+ if (call->op != call_tir_op) {
+ return expr;
+ }
+
+ StructInfo out_sinfo = call->sinfo_args[0];
+ StructInfo infered_sinfo = infer_struct_info_map[op](visited_call,
builder_);
Review Comment:
seems we are mixing legalization and device propagation. Can we leverage
structure info before the the post order visit (aka original Call), assuming
the result dont change during legalization? so we don't have to redo inference.
##########
include/tvm/relax/transform.h:
##########
@@ -666,6 +668,31 @@ TVM_DLL Pass RewriteCUDAGraph();
*/
TVM_DLL Pass FewShotTuning(int valid_count, bool benchmark);
+/*!
+ * \brief This pass is designed to annotate the memory scope information via
VDevice attribute.
+ * This pass need operator attrbutes which in general vanish aftre
legalization.
+ * FuseOps and FuseTIR are modified to pass on the operator specific
attributes and also
+ * op_pattern details as part of the PrimFunc. This pass is Adreno specific
and annotates each
+ * BindingVar with appropriate HintInDevice. RealizeVDevice pass followed by
handles these hints.
+ * Followed by this pass we also invoke SpecializePrimFuncBasedOnCallSite
which updates the
+ * var_buffer_map based on this new VDevice information.
+ */
+TVM_DLL Pass AnnotateCustomMemoryScope(Target target);
+
+/*!
+ * \brief This pass updates the var_buffer mapping of PrimFunctions from the
call_tir info.
+ * Primarily used to update the VDevice information if any changes occured
from the caller.
+ * This pass recreates the buffers and updates the map.
+ */
+TVM_DLL Pass SpecializePrimFuncBasedOnCallSite();
+
+/*
+ * \brief This is a texture specific pass that can optimize unnecessary
to_device copies.
+ * Like texture_scope -> ToVDevice -> global scope. In this case the producer
can directly
+ * store into global scope avoiding unnecessary device copy.
+ */
+TVM_DLL Pass FoldVDeviceScopeChange();
Review Comment:
consider do namespace for this as well, relax/backend/adreno/transform.h
--
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]
