altanh commented on code in PR #11208: URL: https://github.com/apache/tvm/pull/11208#discussion_r883073357
########## 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 Review Comment: would it be more correct to say "minimum required memory" or something along those lines? I'm imagining a scenario where the primitive function internally allocates some buffers (and so the actual memory usage is higher than "live in + live out") ########## 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); + const transform::VarSet& live_out = liveness_.live_out.at(cfg_node); + live_tensors.insert(live_out.begin(), live_out.end()); + live_tensors.erase(call_op); + + // Calculate size of live tensors and store to allow annotation when the function + // gets visited. + uint64_t used_memory = 0; + for (const auto& var : live_tensors) { + Type type = var->checked_type(); + ICHECK(type.defined()) << "InferType pass should be run before AnnotateUsedMemory."; + used_memory += CalculateRelayExprSizeBytes(type); Review Comment: do we have a way of gracefully aborting this analysis if any shapes are dynamic? maybe a guard higher up where this pass gets applied? ########## 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: since the live-in and live-out sets are calculated assuming a fixed/linearized execution order (induced by the input ANF), I think it's possible that the used memory calculation could end up being non-deterministic if ANF linearization is non-deterministic. I'm actually not sure if that's the case with our ANF implementation though, but just something worth keeping in mind. I'm mainly thinking in the case of e.g. ``` y = f(x) z0 = g0(y) ... zn = gn(y) ``` where the `zi`s could really be in any permutation depending on how they get linearized. could you check if the unit tests cover this? ########## 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); + const transform::VarSet& live_out = liveness_.live_out.at(cfg_node); + live_tensors.insert(live_out.begin(), live_out.end()); + live_tensors.erase(call_op); + + // Calculate size of live tensors and store to allow annotation when the function + // gets visited. + uint64_t used_memory = 0; + for (const auto& var : live_tensors) { + Type type = var->checked_type(); + ICHECK(type.defined()) << "InferType pass should be run before AnnotateUsedMemory."; + used_memory += CalculateRelayExprSizeBytes(type); + } + used_memory_annotations_[call_op] = used_memory; Review Comment: I think the IR form at this point prevents this, but I'd add a check that call_op has not already been annotated. If I understand correctly, this pass assumes each primitive func is used at most once. -- 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]
