altanh commented on code in PR #11208: URL: https://github.com/apache/tvm/pull/11208#discussion_r886196916
########## src/relay/backend/annotate_used_memory.cc: ########## @@ -0,0 +1,222 @@ +/* + * 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. + */ + +/*! + * \file src/relay/backend/annotate_used_memory.cc + * \brief Analyzes the used memory at the callsite of primitive functions. + */ + +#include <tvm/ir/module.h> +#include <tvm/relay/attrs/memory.h> +#include <tvm/relay/transform.h> + +#include <unordered_map> +#include <unordered_set> + +#include "../transforms/device_aware_visitors.h" +#include "./liveness_analysis.h" +#include "./utils.h" + +namespace tvm { +namespace relay { +namespace backend { + +/*! + * \brief Annotates the memory usage of each primitive function by analyzing the liveness + * of the input/output tensors at each function callsite and calculating the total amount of + * memory these tensors require. This is added as a "used_memory" annotation to the function + * in question. In addition, the containing function is annotated with an "io_used_memory" + * annotation which refers to the total memory required for the IO tensors. + * + * A simple example: + * + * Before: + * def @main(%input: Tensor[(1, 2, 2, 4), int8]) -> Tensor[(1, 2, 2, 4), int8] { + * let %x_0 = fn (%x: Tensor[(1, 2, 2, 4), int8], Primitive=1) -> Tensor[(1, 2, 2, 4), int8] { + * nn.max_pool2d(%x, pool_size=[1, 1], padding=[0, 0, 0, 0]) + * }; + * let %x_1 = %x_0(%input); + * %x_1 + * } + * + * After: + * def @main(%input: Tensor[(1, 2, 2, 4), int8], io_used_memory=32) -> Tensor[(1, 2, 2, 4), int8] { + * let %x_0: fn (%x: Tensor[(1, 2, 2, 4), int8], Primitive=1, used_memory=32) -> Tensor[(1, 2, 2, + * 4), int8] { nn.max_pool2d(%x, pool_size=[1, 1], padding=[0, 0, 0, 0]) + * }; + * let %x_1: Tensor[(1, 2, 2, 4), int8] = %x_0(%input); + * %x_1 + * } + * + * Note that in the simple example above io_used_memory and used_memory are the same since there + * is only one primitive function. + */ +class AnnotateUsedMemoryMutator : public transform::DeviceAwareExprMutator { + public: + AnnotateUsedMemoryMutator(const IRModule& module, const transform::ControlFlowGraph& cfg, + const transform::LivenessAnalysis& lva) + : DeviceAwareExprMutator(module), control_flow_graph_(cfg), liveness_(lva) {} + + /*! + * \brief Mutates the input function. In addition, an "io_used_memory" annotation is + * added to the input function which refers to the total size required for the IO + * tensors. + */ + Function operator()(const Function& func) { + uint64_t io_used_memory = 0; + + // Inputs + for (const Var& param : func->params) { + Type type = param->checked_type(); + ICHECK(type.defined()) << "InferType pass should be run before AnnotateUsedMemory."; + io_used_memory += CalculateRelayExprSizeBytes(type); + } + + // Outputs + Type type = func->body->checked_type(); + ICHECK(type.defined()) << "InferType pass should be run before AnnotateUsedMemory."; + io_used_memory += CalculateRelayExprSizeBytes(type); + + Expr new_func_body = VisitExpr(func->body); + Function new_func = WithFields(func, func->params, new_func_body); + return WithAttr(std::move(new_func), "io_used_memory", + tvm::IntImm(tvm::DataType::UInt(64), io_used_memory)); + } + + /*! + * \brief Establish which let bindings have primitive function values. + */ + std::pair<Var, Expr> PreVisitLetBinding_(const Var& var, const Expr& value) { + if (const auto* func_node = value.as<FunctionNode>()) { + ICHECK(func_node->attrs.HasNonzeroAttr(attr::kPrimitive)) + << "Expect top-level functions to be primitive."; + let_bound_prim_func_.insert(var); + } + return DeviceAwareExprMutator::PreVisitLetBinding_(var, value); + } + + /*! + * \brief Visit let nodes and perform one of two actions depending on their value: + * + * 1. CallNode - Calculate "used_memory" annotation value at the callsite of + * primitive functions. + * + * 2. FunctionNode - Annotate functions with "used_memory" annotation based on the + * previous analysis at the callsite. + * + */ + Expr PostVisitLet_(const LetNode* pre_let_node, const LetNode* post_let_node) override { + Var let_var = post_let_node->var; + Expr let_value = IgnoreOnDevice(post_let_node->value); + + if (let_value->IsInstance<CallNode>()) { + Call callsite = Downcast<Call>(let_value); + if (CheckPrimitiveFunctionCall(callsite)) { + Var call_op = Downcast<Var>(callsite->op); + + // Find all the vars that are live at the callsite. This is done by merging the + // in and out varset's and then removing the var that references the primitive + // function itself since we don't want this included in the calculation. + const transform::ControlFlowGraph::NodePtr cfg_node = + control_flow_graph_.let_map.at(GetRef<Let>(pre_let_node)); + transform::VarSet live_tensors = liveness_.live_in.at(cfg_node); Review Comment: kind of, I think you've got my point- my comment was mainly about how the same function might get different used-mem annotations depending on the incoming ANF ordering. In the unit test, you only rely on the ordering to match up the used mems (which will be correct), but there could be some trickiness if we expect the same function (say, referenced by name or id) to have the same annotation between runs of the pass (*if* ANF ordering is non-deterministic). Let me know if that makes sense! This might not be a problem for your use case anyways. -- 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]
