junrushao1994 commented on code in PR #12019:
URL: https://github.com/apache/tvm/pull/12019#discussion_r915288059
##########
src/tir/transforms/compact_buffer_region.cc:
##########
@@ -45,25 +45,43 @@ using support::NDIntSet;
* \brief simplify and return the region collected by NDIntSet. return the
original
* buffer shape if the int_set is empty.
*/
-Region SimplifyAndNarrowBufferRegionFromNDIntSet(const NDIntSet& nd_int_set,
- const Array<PrimExpr>&
original_shape,
- arith::Analyzer* analyzer) {
+Region SimplifyAndNarrowBufferRegionFromNDIntSet(
+ const NDIntSet& nd_int_set, const Array<PrimExpr>& original_shape,
arith::Analyzer* analyzer,
+ const std::vector<const ForNode*>& ancestor_loops) {
Array<Range> result;
result.reserve(nd_int_set.size());
for (size_t i = 0; i < nd_int_set.size(); ++i) {
const arith::IntSet& int_set = nd_int_set[i];
Range range = int_set.CoverRange(Range(/*begin=*/0,
/*end=*/original_shape[i]));
- result.push_back(
- Range::FromMinExtent(analyzer->Simplify(max(0, range->min)),
- analyzer->Simplify(min(original_shape[i],
range->extent))));
+ PrimExpr min = analyzer->Simplify(tvm::max(0, range->min));
+ PrimExpr extent = analyzer->Simplify(tvm::min(original_shape[i],
range->extent));
+
+ // Check the buffer region is not loop dependent, since loop dependent
+ // allocation is not supported yet.
+ auto is_loop_var = [&ancestor_loops](const VarNode* v) {
+ return std::any_of(ancestor_loops.begin(), ancestor_loops.end(),
+ [v](const ForNode* n) { return n->loop_var.get() ==
v; });
+ };
+ if (UsesVar(extent, is_loop_var)) {
+ // try estimate a constant upperbound on region's extent
+ int64_t upperbound = analyzer->const_int_bound(extent)->max_value;
+ if (upperbound != arith::ConstIntBound::kPosInf) {
+ extent = make_const(extent->dtype, upperbound);
+ } else {
+ // or else we have to fallback to full region
+ min = 0;
Review Comment:
use `make_zero(original_shape[i]->dtype)`
--
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]
