From 958ca0cc3698a087229a51511adc8708023201fa Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Thu, 11 Apr 2024 15:41:10 -0400 Subject: [PATCH] [Relax] Enhance symbolic expr estimation in memory planning This PR enhances the symbolic expression upper bound estimation in static memory planning. Prior to this PR, we are not able to estimate the upper bound of `a * b` when `a` has an upper bound while `b` does not. This PR enhances the estimation with arith::IntSet. We introduce another TIR attribute `tir_non_negative_var` to indicate the non-negative TIR variables for memory planning use. A new unit test is introduced for this enhancement. --- .../transform/static_plan_block_memory.cc | 45 ++++++-- ...test_transform_static_plan_block_memory.py | 102 ++++++++++++++++++ 2 files changed, 137 insertions(+), 10 deletions(-) diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc index 453c99691613..2b16d8650906 100644 --- a/src/relax/transform/static_plan_block_memory.cc +++ b/src/relax/transform/static_plan_block_memory.cc @@ -353,8 +353,10 @@ class StorageAllocatorBaseVisitor : public ExprVisitor { * the input function signature in the analyzer. * \param func The function to be analyzed. * \param ana The analyzer which contains the TIR var upper bounds. + * \param dom_map The domain map of the TIR variables. */ -void SetTIRVarUpperBound(Function func, arith::Analyzer* ana) { +void SetTIRVarUpperBound(Function func, arith::Analyzer* ana, + Map* dom_map) { // Use the attribute-annotated TIR var upper bounds as the TIR var values for // memory planning. // NOTE: we only apply the annotated upper bounds to the TIR variables that @@ -362,7 +364,10 @@ void SetTIRVarUpperBound(Function func, arith::Analyzer* ana) { Map var_upper_bound_attr_raw = func->GetAttr>("tir_var_upper_bound") .value_or(Map()); + Array non_negative_var_attr_raw = + func->GetAttr>("tir_non_negative_var").value_or(Array()); std::unordered_map var_upper_bound_attr; + std::unordered_set non_negative_var_attr; // We manually check the value type to ensure the values are all positive IntImm. for (auto it : var_upper_bound_attr_raw) { const auto* key = it.first.as(); @@ -378,13 +383,23 @@ void SetTIRVarUpperBound(Function func, arith::Analyzer* ana) { << value->value << " is got."; var_upper_bound_attr[GetRef(key)] = GetRef(value); } + for (ObjectRef var_name : non_negative_var_attr_raw) { + const auto* key = var_name.as(); + CHECK(key != nullptr) << "The element of attr `tir_non_negative_var` should be string. However " + << key->GetTypeKey() << " is got."; + non_negative_var_attr.insert(GetRef(key)); + } Array var_in_signature = TIRVarsInStructInfo(GetStructInfo(func)); for (const tir::Var& tir_var : var_in_signature) { auto it = var_upper_bound_attr.find(tir_var->name_hint); if (it != var_upper_bound_attr.end()) { - ana->Bind(tir_var, - tvm::Range::FromMinExtent(tvm::IntImm(DataType::Int(64), 0), - tvm::IntImm(DataType::Int(64), (*it).second->value + 1))); + tvm::Range range = + tvm::Range::FromMinExtent(tvm::IntImm(DataType::Int(64), 0), + tvm::IntImm(DataType::Int(64), (*it).second->value + 1)); + ana->Bind(tir_var, range); + dom_map->Set(tir_var, arith::IntSet::FromRange(range)); + } else if (non_negative_var_attr.count(tir_var->name_hint)) { + ana->MarkGlobalNonNegValue(tir_var); } } } @@ -398,14 +413,20 @@ void SetTIRVarUpperBound(Function func, arith::Analyzer* ana) { * \return The upper-bounded shape. When a dimension's upper bound * cannot be determined, we keep the dimension unchanged. */ -Array GetUpperBoundShape(Array shape, arith::Analyzer* ana) { +Array GetUpperBoundShape(Array shape, arith::Analyzer* ana, + const Map& dom_map) { // Use the upper bounds of TIR vars as their values. Array upper_bounded_shape; upper_bounded_shape.reserve(shape.size()); for (const PrimExpr& dim_len : shape) { int64_t max_bound = ana->const_int_bound(dim_len)->max_value; if (max_bound == std::numeric_limits::max()) { - upper_bounded_shape.push_back(dim_len); + arith::IntSet int_set = ana->int_set(dim_len, dom_map); + if (int_set.HasUpperBound()) { + upper_bounded_shape.push_back(int_set.max()); + } else { + upper_bounded_shape.push_back(dim_len); + } } else { upper_bounded_shape.push_back(tvm::IntImm(DataType::Int(64), max_bound)); } @@ -462,7 +483,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { void VisitExpr_(const FunctionNode* func) final { // Set the upper bound of TIR variables in the analyzer. - SetTIRVarUpperBound(GetRef(func), analyzer_); + SetTIRVarUpperBound(GetRef(func), analyzer_, &dom_map_); // Recurse into the function to get its tokens. Tokens body_tokens = GetTokens(func->body); // Discard the tokens used by the function return value, as they are external referenced. @@ -565,7 +586,7 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { // Use the upper bounds of TIR vars as their values. The upper bound shape can still be dynamic // if the upper bounds of some variables are not provided. - Array upper_bounded_shape = GetUpperBoundShape(shape->values, analyzer_); + Array upper_bounded_shape = GetUpperBoundShape(shape->values, analyzer_, dom_map_); // Create and set token. StringImm storage_scope = Downcast(call->args[3]); @@ -641,6 +662,8 @@ class StorageAllocatorInit : public StorageAllocatorBaseVisitor { const IRModule& ctx_mod_; /*! \brief The arithmetic analyzer. */ arith::Analyzer* analyzer_; + /*! \brief The domain map of dynamic TIR variables for analysis. */ + Map dom_map_; /*! \brief The mapping from each token to the binding block where it is created. */ std::unordered_map token2block_; /*! \brief The mapping from each token to the Exprs that are using this token. */ @@ -816,7 +839,7 @@ class StorageAllocationRewriter : public ExprMutator { plan_dynamic_output_ = static_cast( func_->GetAttr(plan_dyn_attr_).value_or(IntImm(DataType::Int(32), 0))->value); if (plan_dynamic_output_) { - SetTIRVarUpperBound(GetRef(func_), &ana_); + SetTIRVarUpperBound(GetRef(func_), &ana_, &dom_map_); } token2storage_var_.clear(); Function func = Downcast(this->VisitExpr_(func_)); @@ -879,7 +902,7 @@ class StorageAllocationRewriter : public ExprMutator { ICHECK_NOTNULL(sinfo); const auto* shape = sinfo->shape.as(); ICHECK_NOTNULL(shape); - Array upper_bounded_shape = GetUpperBoundShape(shape->values, &ana_); + Array upper_bounded_shape = GetUpperBoundShape(shape->values, &ana_, dom_map_); if (!IsStaticShape(shape->values)) { ICHECK(!sinfo->IsUnknownDtype()); ICHECK_EQ(sinfo->dtype, Downcast(call->args[1])->value); @@ -906,6 +929,8 @@ class StorageAllocationRewriter : public ExprMutator { /*! \brief The arithmetic analyzer. */ arith::Analyzer ana_; + /*! \brief The domain map of dynamic TIR variables for analysis. */ + Map dom_map_; /*! \brief A boolean indicating whether to plan dynamic-shape function output tensors. */ bool plan_dynamic_output_; /*! diff --git a/tests/python/relax/test_transform_static_plan_block_memory.py b/tests/python/relax/test_transform_static_plan_block_memory.py index 83eff854a40d..63f422d4cfbe 100644 --- a/tests/python/relax/test_transform_static_plan_block_memory.py +++ b/tests/python/relax/test_transform_static_plan_block_memory.py @@ -1347,5 +1347,107 @@ def main(x: R.Tensor((2, "n"), dtype="float32")): relax.transform.StaticPlanBlockMemory()(Module) +def test_add(): + @I.ir_module + class Module: + @T.prim_func(private=True) + def cumsum(var_A: T.handle, var_A_1: T.handle, var_exclusive_scan_thrust: T.handle): + T.evaluate(0) + + @R.function + def main( + probs: R.Tensor(("batch_size", "vocab_size"), dtype="float32") + ) -> R.Tensor(("batch_size", "vocab_size"), dtype="float32"): + batch_size = T.int64() + vocab_size = T.int64() + R.func_attr( + { + "relax.force_pure": 1, + "relax.memory_plan_dynamic_func_output": 1, + "tir_var_upper_bound": {"batch_size": 32}, + "tir_non_negative_var": ["vocab_size"], + } + ) + cls = Module + lv1: R.Tensor( + (2 * (batch_size * vocab_size * 4) + 4194304,), + dtype="uint8", + ) = R.builtin.alloc_tensor( + R.shape([2 * (batch_size * vocab_size * 4) + 4194304]), + R.dtype("uint8"), + R.prim_value(0), + R.str("global"), + ) + alloc1: R.Tensor((batch_size, vocab_size), dtype="float32") = R.builtin.alloc_tensor( + R.shape([batch_size, vocab_size]), + R.dtype("float32"), + R.prim_value(0), + R.str("global"), + ) + cls.cumsum(probs, lv1, alloc1) + cumsum: R.Tensor((batch_size, vocab_size), dtype="float32") = alloc1 + lv1_1: R.Tensor((batch_size, vocab_size), dtype="int32") = R.call_packed( + "vm.builtin.reshape", + cumsum, + R.shape([batch_size, vocab_size]), + sinfo_args=(R.Tensor((batch_size, vocab_size), dtype="float"),), + ) + return lv1_1 + + @I.ir_module + class Expected: + @T.prim_func(private=True) + def cumsum(var_A: T.handle, var_A_1: T.handle, var_exclusive_scan_thrust: T.handle): + T.evaluate(0) + + @R.function + def main( + probs: R.Tensor(("batch_size", "vocab_size"), dtype="float32") + ) -> R.Tensor(("batch_size", "vocab_size"), dtype="int32"): + batch_size = T.int64() + vocab_size = T.int64() + R.func_attr( + { + "relax.force_pure": 1, + "tir_non_negative_var": ["vocab_size"], + "tir_var_upper_bound": {"batch_size": 32}, + } + ) + cls = Expected + storage: R.Object = R.memory.alloc_storage( + R.shape([32 * vocab_size * 4 * 2 + 4194304]), + R.prim_value(0), + R.str("global"), + R.dtype("uint8"), + ) + lv1: R.Tensor( + (2 * (batch_size * vocab_size * 4) + 4194304,), + dtype="uint8", + ) = R.memory.alloc_tensor( + storage, + R.prim_value(0), + R.shape([2 * (batch_size * vocab_size * 4) + 4194304]), + R.dtype("uint8"), + ) + storage1: R.Object = R.memory.alloc_storage( + R.shape([128 * vocab_size]), R.prim_value(0), R.str("global"), R.dtype("float32") + ) + alloc1: R.Tensor((batch_size, vocab_size), dtype="float32") = R.memory.alloc_tensor( + storage1, R.prim_value(0), R.shape([batch_size, vocab_size]), R.dtype("float32") + ) + cls.cumsum(probs, lv1, alloc1) + cumsum: R.Tensor((batch_size, vocab_size), dtype="float32") = alloc1 + lv1_1: R.Tensor((batch_size, vocab_size), dtype="int32") = R.call_packed( + "vm.builtin.reshape", + cumsum, + R.shape([batch_size, vocab_size]), + sinfo_args=(R.Tensor((batch_size, vocab_size), dtype="float32"),), + ) + return lv1_1 + + mod = relax.transform.StaticPlanBlockMemory()(Module) + tvm.ir.assert_structural_equal(mod, Expected) + + if __name__ == "__main__": tvm.testing.main()