diff --git a/include/tvm/tir/buffer.h b/include/tvm/tir/buffer.h index 28d202cb50a9..f04209d0b061 100644 --- a/include/tvm/tir/buffer.h +++ b/include/tvm/tir/buffer.h @@ -121,6 +121,14 @@ class BufferNode : public Object { return shape.size() != 0 ? shape[0].dtype() : DataType::Int(32); } + /*! \brief Determine the offset in the buffer of the given index. + * + * Returns the buffer offset, in number of elements of type dtype, + * without adjusting for number of lanes. (e.g. The number of + * float16x4 elements in a buffer of type float16x4.) + */ + PrimExpr ElemOffset(Array index) const; + static constexpr const char* _type_key = "tir.Buffer"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; diff --git a/src/tir/ir/buffer.cc b/src/tir/ir/buffer.cc index 8253fa3c0a36..24aacc3c04f7 100644 --- a/src/tir/ir/buffer.cc +++ b/src/tir/ir/buffer.cc @@ -246,41 +246,41 @@ inline PrimExpr MergeMulMod(arith::Analyzer* analyzer, const PrimExpr& base) { // The buffer offset in convention of number of elements of // original data ignoring number of lanes. // We also perform optimization to simplify the indexing expression. -inline PrimExpr ElemOffset(const BufferNode* n, Array index) { - PrimExpr base = n->elem_offset; +PrimExpr BufferNode::ElemOffset(Array index) const { + PrimExpr base = this->elem_offset; arith::Analyzer ana; - if (n->strides.size() == 0) { + if (this->strides.size() == 0) { // Scalar case - if (n->shape.size() == 0 && index.size() == 1) { + if (this->shape.size() == 0 && index.size() == 1) { auto is_int = index[0].as(); ICHECK(is_int && is_int->value == 0); base = base + index[0]; } else { - ICHECK_EQ(n->shape.size(), index.size()); + ICHECK_EQ(this->shape.size(), index.size()); if (index.size() > 0) { PrimExpr offset = index[0]; for (size_t i = 1; i < index.size(); ++i) { - offset = MergeMulMod(&ana, offset * n->shape[i] + index[i]); + offset = MergeMulMod(&ana, offset * this->shape[i] + index[i]); } base = base + offset; } } } else { - ICHECK_EQ(n->strides.size(), index.size()); + ICHECK_EQ(this->strides.size(), index.size()); if (is_zero(base)) { - base = MergeMulMod(&ana, index[0] * n->strides[0]); + base = MergeMulMod(&ana, index[0] * this->strides[0]); } else { - base = MergeMulMod(&ana, base + index[0] * n->strides[0]); + base = MergeMulMod(&ana, base + index[0] * this->strides[0]); } for (size_t i = 1; i < index.size(); ++i) { - base = MergeMulMod(&ana, base + index[i] * n->strides[i]); + base = MergeMulMod(&ana, base + index[i] * this->strides[i]); } } return base; } inline PrimExpr BufferOffset(const BufferNode* n, Array index, DataType dtype) { - PrimExpr offset = ElemOffset(n, index); + PrimExpr offset = n->ElemOffset(index); if (n->dtype.lanes() != 1) { offset = offset * make_const(offset.dtype(), dtype.lanes()); } @@ -353,7 +353,7 @@ Buffer Buffer::MakeSlice(Array begins, Array extents) const ICHECK(n != nullptr); arith::Analyzer ana; begins = SimplifyArray(&ana, begins); - PrimExpr elem_offset = ana.Simplify(ElemOffset(n, begins)); + PrimExpr elem_offset = ana.Simplify(n->ElemOffset(begins)); Array strides = n->strides; if (strides.size() == 0) { bool can_relax = true; diff --git a/src/tir/transforms/arg_binder.cc b/src/tir/transforms/arg_binder.cc index 293c990d2745..d3ab32cbd7f9 100644 --- a/src/tir/transforms/arg_binder.cc +++ b/src/tir/transforms/arg_binder.cc @@ -204,7 +204,7 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type, def_handle_dtype_.Set(v_strides, tir::TypeAnnotation(tvm_shape_type)); init_nest_.emplace_back( LetStmt(v_strides, TVMArrayGet(DataType::Handle(), handle, builtin::kArrStrides), nop)); - PrimExpr is_null = Call(DataType::Bool(1), builtin::isnullptr(), {v_strides}); + PrimExpr v_strides_is_null = Call(DataType::Bool(1), builtin::isnullptr(), {v_strides}); if (buffer->strides.size() == 0) { // Assert the buffer is compact DataType stype = buffer->DefaultIndexType(); @@ -226,7 +226,7 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type, foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_and(a, b, span); }, const_true(1), conds), stride_msg, Evaluate(0)); - check = IfThenElse(Not(is_null), check, Stmt()); + check = IfThenElse(Not(v_strides_is_null), check, Stmt()); asserts_.emplace_back(SeqStmt({check, Evaluate(0)})); } } else if (buffer->buffer_type == kAutoBroadcast) { @@ -239,24 +239,29 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type, PrimExpr value = cast(buffer->shape[k].dtype(), Load(tvm_shape_type, v_strides, IntImm(DataType::Int(32), k), const_true(1))); - value = tvm::if_then_else(is_null, stride, value); + value = tvm::if_then_else(v_strides_is_null, stride, value); value = tvm::if_then_else(buffer->shape[k] == 1, 0, value); Bind_(buffer->strides[k], value, field_name.str(), true); stride = analyzer_.Simplify(stride * buffer->shape[k]); } } else { - std::ostringstream stride_null_err_msg; - stride_null_err_msg << arg_name << ".strides: expected non-null strides."; - asserts_.emplace_back( - AssertStmt(Not(is_null), tvm::tir::StringImm(stride_null_err_msg.str()), nop)); + PrimExpr stride_from_shape = 1; - for (size_t k = 0; k < buffer->strides.size(); ++k) { + for (int k = buffer->strides.size() - 1; k >= 0; k--) { std::ostringstream field_name; field_name << v_strides->name_hint << '[' << k << ']'; + + PrimExpr explicit_stride = + cast(buffer->shape[k].dtype(), + Load(tvm_shape_type, v_strides, IntImm(DataType::Int(32), k), const_true(1))); + Bind_(buffer->strides[k], - cast(buffer->shape[k].dtype(), - Load(tvm_shape_type, v_strides, IntImm(DataType::Int(32), k), const_true(1))), + tvm::if_then_else(v_strides_is_null, stride_from_shape, explicit_stride), field_name.str(), true); + + stride_from_shape *= + cast(buffer->shape[k].dtype(), + Load(tvm_shape_type, v_shape, IntImm(DataType::Int(32), k), const_true(1))); } } // Byte_offset field. diff --git a/src/tir/transforms/storage_flatten.cc b/src/tir/transforms/storage_flatten.cc index 2c32cc7f0883..6a3ce596c2fe 100644 --- a/src/tir/transforms/storage_flatten.cc +++ b/src/tir/transforms/storage_flatten.cc @@ -50,6 +50,984 @@ using runtime::StorageRank; using runtime::StorageScope; using runtime::ThreadScope; +/* Make buffer realize extents and buffer shapes consistent + * + * For external buffers, verify that the extents of BufferRealize + * nodes match the shape of the external buffer. For internal + * buffers, rewrite the shape of the Buffer objects to match the + * extent of the BufferRealize, and rewrite indices of + * BufferLoad/BufferStore nodes to match. + */ +class BufferShapeLegalize : public StmtExprMutator { + public: + explicit BufferShapeLegalize(const Map& extern_buffer_map, + IRVisitorWithAnalyzer* bound_analyzer) + : bound_analyzer_(bound_analyzer) { + for (auto kv : extern_buffer_map) { + Buffer buf = kv.second; + extern_buffers_.insert(buf); + + BufferEntry remap; + remap.remap_to = buf; + remap.index_offsets = Array(buf->shape.size(), 0); + remap.in_scope = true; + buf_map_[buf] = remap; + } + } + + PrimExpr VisitExpr_(const VarNode* op) final { + auto it = var_remap_.find(op); + if (it != var_remap_.end()) { + return it->second; + } else { + return GetRef(op); + } + } + + Stmt VisitStmt_(const BufferRealizeNode* op) final { + // BufferRealizeNode for an external buffer serves as an + // annotation of the external buffers, and should not be changed. + // Instead, verify that the bounds match the external + // buffer. + if (extern_buffers_.count(op->buffer)) { + CHECK_EQ(op->buffer->shape.size(), op->bounds.size()) + << "External buffer realize has mismatched dimension"; + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + ICHECK(op); + + for (size_t i = 0; i < op->bounds.size(); i++) { + PrimExpr eq = bound_analyzer_->Simplify(op->buffer->shape[i] == op->bounds[i]->extent); + std::ostringstream ss; + ss << "Dim " << i << " of external buffer " << op->buffer->name << " has shape " + << op->buffer->shape[i] << ", but is only realized for extent " << op->bounds[i]->extent; + if (auto eq_int = eq.as()) { + ICHECK(eq_int->value) << ss.str(); + } else { + stmt = AssertStmt(eq, tvm::tir::StringImm(ss.str()), stmt); + } + } + return stmt; + } + + // Compute the new buffer shape, new realization bounds, and the + // offsets to be applied to buffer access. + Array realized_shape; + Array index_offsets; + Array new_bounds; + for (size_t i = 0; i < op->bounds.size(); i++) { + const Range& bound = op->bounds[i]; + realized_shape.push_back(bound->extent); + index_offsets.push_back(bound->min); + new_bounds.push_back({0, bound->extent}); + } + + if (op->buffer->shape.size()) { + ICHECK_EQ(op->buffer->shape.size(), realized_shape.size()) + << "Inconsistency between dimension of buffer " << op->buffer + << " and dimension of its realized bounds."; + } + + Buffer key = op->buffer; + + Buffer buf = op->buffer; + auto write_ptr = buf.CopyOnWrite(); + write_ptr->shape = realized_shape; + + { + BufferEntry remap; + remap.remap_to = buf; + remap.index_offsets = index_offsets; + remap.in_scope = true; + buf_map_[key] = remap; + } + + Stmt stmt = BufferRealize(buf, new_bounds, op->condition, this->VisitStmt(op->body), op->span); + + buf_map_.at(key).in_scope = false; + + return stmt; + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + ICHECK(op); + + auto it = buf_map_.find(op->buffer); + if (it != buf_map_.end()) { + const BufferEntry& entry = it->second; + ICHECK(entry.in_scope) << "Cannot store to an out-of-scope buffer"; + + BufferStore updated = GetRef(op); + auto write_ptr = updated.CopyOnWrite(); + write_ptr->indices = update_indices(op->indices, entry.index_offsets); + write_ptr->buffer = entry.remap_to; + stmt = updated; + } + + return stmt; + } + + PrimExpr VisitExpr_(const BufferLoadNode* op) final { + PrimExpr expr = StmtExprMutator::VisitExpr_(op); + op = expr.as(); + ICHECK(op); + + auto it = buf_map_.find(op->buffer); + if (it != buf_map_.end()) { + const BufferEntry& entry = it->second; + ICHECK(entry.in_scope) << "Cannot read from an out-of-scope buffer"; + + BufferLoad updated = GetRef(op); + auto write_ptr = updated.CopyOnWrite(); + write_ptr->indices = update_indices(op->indices, entry.index_offsets); + write_ptr->buffer = entry.remap_to; + expr = updated; + } + + return expr; + } + + Stmt VisitStmt_(const AttrStmtNode* op) final { + if (op->node->IsInstance()) { + // Visit body before checking internal_buf_map_, because we + // don't know if the BufferNode needs to be changed until we + // look in the body for a BufferRealizeNode with different + // extents. + Stmt body = this->VisitStmt(op->body); + + Buffer buffer = Downcast(op->node); + auto it = buf_map_.find(buffer); + if (it != buf_map_.end()) { + buffer = it->second.remap_to; + return AttrStmt(it->second.remap_to, op->attr_key, op->value, body); + } + return AttrStmt(buffer, op->attr_key, op->value, body); + + } else if (op->attr_key == attr::buffer_bind_scope) { + return HandleBufferBindScope(op); + } + + return StmtExprMutator::VisitStmt_(op); + } + + private: + // Any buffers that give views into a resized buffer should be + // updated, both to refer to the resized buffer and to have the view + // window updated. For example, suppose B1 is a 1-D buffer of size + // 100 which is only realized on the range (10,50), and buffer V1 is + // a view into B1[25:35]. When B1 is replaced with B2, a buffer of + // size 40 realized on the range (0,40), V1 must be replaced to be a + // view into B2[15:25]. + Stmt HandleBufferBindScope(const AttrStmtNode* op) { + Array arr = Downcast>(op->node); + ICHECK_EQ(arr.size(), 2U); + Buffer buffer = Downcast(arr[0]); + ICHECK(buffer.defined()); + Buffer target = Downcast(arr[1]); + ICHECK(target.defined()); + + auto it = buf_map_.find(target); + ICHECK(it != buf_map_.end()) << "attr::buffer_bind_scope target " << target << " not in scope."; + const BufferEntry& target_remap = it->second; + + ICHECK(target_remap.in_scope) << "Cannot bind " << buffer->name + << " to the out-of-scope buffer " << target_remap.remap_to->name; + + Call tuple = Downcast(op->value); + ICHECK(tuple.defined() && tuple->op.same_as(builtin::tvm_tuple())); + + Array new_tuple_args; + Array realized_begins; + Array view_shape; + ICHECK_EQ(tuple->args.size(), target_remap.index_offsets.size() * 2) + << "attr::buffer_bind_scope to define " << buffer << " as a view into " << target + << " does match dimensionality of " << target; + for (size_t i = 0; i < target_remap.index_offsets.size(); i++) { + PrimExpr parent_begin = tuple->args[2 * i]; + PrimExpr view_extent = tuple->args[2 * i + 1]; + // Offset the begin of the buffer view by the offset of the target buffer. + new_tuple_args.push_back(parent_begin - target_remap.index_offsets[i]); + // Keep the extent of the buffer view the same. + new_tuple_args.push_back(view_extent); + // Use the extent of the buffer view to define the buffer view's shape. + view_shape.push_back(view_extent); + // Within the buffer view, indices start at 0. + realized_begins.push_back(0); + } + + // If a view is binding to a buffer of a higher dimensionality, + // then the leading dimensions should be padded out with shape of + // 1. + ICHECK_GE(view_shape.size(), buffer->shape.size()) + << "Cannot bind " << buffer << " to a shape of lower dimension."; + if (view_shape.size() > buffer->shape.size()) { + size_t diff = view_shape.size() - buffer->shape.size(); + Array padded_shape; + for (size_t i = 0; i < diff; i++) { + padded_shape.push_back(1); + } + for (auto dim : buffer->shape) { + padded_shape.push_back(dim); + } + view_shape = std::move(padded_shape); + } + + // If a buffer has strides defined, and is being remapped into a + // shape with additional dimensions, then define dummy values for + // the strides. + Array realized_strides = buffer->strides; + if ((realized_strides.size() > 0) && (realized_strides.size() != view_shape.size())) { + ICHECK_GE(view_shape.size(), realized_strides.size()) + << "Cannot bind the strides of " << buffer << " to a shape of lower dimension"; + size_t diff = view_shape.size() - buffer->strides.size(); + + Array updated_strides; + for (size_t i = 0; i < diff; i++) { + updated_strides.push_back(Var("stride", buffer->shape[0].dtype())); + } + for (auto stride : buffer->strides) { + updated_strides.push_back(stride); + } + realized_strides = updated_strides; + } + + Buffer key = buffer; + + auto write_ptr = buffer.CopyOnWrite(); + write_ptr->shape = view_shape; + write_ptr->strides = realized_strides; + + { + BufferEntry remap; + remap.index_offsets = realized_begins; + remap.remap_to = buffer; + remap.in_scope = true; + buf_map_[key] = remap; + } + + // Define remappings of any Variables referencing Buffer internals + // (e.g. Store/Load nodes). Passing fuzzy_match=true allows the + // remapped buffer to have a number of dimensions. + ArgBinder binder(&var_remap_); + binder.BindBuffer(key, buffer, key->name, true); + + Stmt body = this->VisitStmt(op->body); + body = MergeNest(binder.asserts(), body); + body = MergeNest(binder.init_nest(), body); + + Stmt stmt = AttrStmt(Array{buffer, target_remap.remap_to}, op->attr_key, + Call(tuple->dtype, tuple->op, new_tuple_args, tuple->span), body); + + for (const Var& v : binder.defs()) { + var_remap_.erase(v.get()); + } + + buf_map_.at(key).in_scope = false; + return stmt; + } + + Array update_indices(const Array& indices, const Array& offsets) { + // offsets come from BufferRealizeNode::bounds, which is allowed + // to be empty to indicate realization of the full shape of the + // buffer. In that case, the indices do not need to be modified, + // but may need to be extended with leading zeroes. + if (offsets.size() == 0) { + return indices; + } + + ICHECK_GE(offsets.size(), indices.size()) + << "Cannot bind buffer to a shape of lower dimension."; + + Array new_indices; + + // Pad leading indices with zero, matching the "fuzzy_match" + // behavior from ArgBinder::BindBuffer. + size_t diff = offsets.size() - indices.size(); + for (size_t i = 0; i < diff; i++) { + new_indices.push_back(0); + } + + // Offset indices used to access buffers of a reduced size. + for (size_t i = 0; i < indices.size(); i++) { + PrimExpr offset = offsets[i + diff]; + new_indices.push_back(indices[i] - offset); + } + + return new_indices; + } + + std::unordered_map var_remap_; + + std::unordered_set extern_buffers_; + + struct BufferEntry { + Buffer remap_to; + Array index_offsets; + bool in_scope; + }; + + std::unordered_map buf_map_; + + IRVisitorWithAnalyzer* bound_analyzer_; +}; + +/* Apply dimension alignment restrictions + * + * Buffers annotated with attr::buffer_dim_align may need to have + * strides defined such that they are no longer in a compact shape. + * After this pass, buffers have stride definitions to include these + * alignment restrictions, and attr::buffer_dim_align annotations have + * been removed. + */ +class BufferStrideLegalize : public StmtExprMutator { + public: + explicit BufferStrideLegalize(const Map& extern_buffer_map, + IRVisitorWithAnalyzer* bound_analyzer) + : bound_analyzer_(bound_analyzer) { + for (auto kv : extern_buffer_map) { + Buffer buf = kv.second; + Buffer with_strides = WithStrides(buf); + { + BufferEntry entry; + entry.remap_to = with_strides; + entry.in_scope = true; + entry.is_external = true; + buf_map_[buf] = entry; + } + updated_extern_buffer_map_.Set(kv.first, with_strides); + } + } + + Map UpdatedExternBufferMap() const { return updated_extern_buffer_map_; } + + Buffer WithStrides(Buffer buf) { + auto it = buf_map_.find(buf); + if (it != buf_map_.end()) { + const BufferEntry& entry = it->second; + ICHECK(entry.in_scope) << "Cannot annotate an out-of-scope buffer"; + return entry.remap_to; + } + + if (buf->strides.size()) { + ICHECK_EQ(buf->strides.size(), buf->shape.size()) + << "Buffer " << buf << " has inconsistent strides/shape."; + return buf; + } + + // Keeping this to have matched behavior to previous version. + // There are many parts of the codebase that assume that a strided + // array cannot be compact. For example, ArgBinder::BindBuffer + // and tir.Specialize. + if (dim_align_.count(buf) == 0) { + return buf; + } + + // Can't define the strides for a buffer without a known shape. + Array shape = buf->shape; + if (shape.size() == 0) { + return buf; + } + + std::vector rstrides; + const std::vector& avec = dim_align_[buf]; + int first_dim = 0; + PrimExpr stride = make_const(shape[first_dim].dtype(), 1); + for (size_t i = shape.size(); i != 0; --i) { + size_t dim = i - 1; + if (dim < avec.size() && avec[dim].align_factor != 0) { + PrimExpr factor = make_const(stride.dtype(), avec[dim].align_factor); + PrimExpr offset = make_const(stride.dtype(), avec[dim].align_offset); + stride = stride + indexmod(factor + offset - indexmod(stride, factor), factor); + stride = bound_analyzer_->Simplify(stride); + } + rstrides.push_back(stride); + stride = stride * shape[dim]; + } + + auto ptr = buf.CopyOnWrite(); + ptr->strides = Array(rstrides.rbegin(), rstrides.rend()); + + return buf; + } + + Stmt VisitStmt_(const AttrStmtNode* op) final { + if (op->attr_key == attr::buffer_dim_align) { + auto buffer = Downcast(op->node); + const CallNode* tuple = op->value.as(); + ICHECK(tuple && tuple->op.same_as(builtin::tvm_tuple())); + auto& vinfo = dim_align_[buffer]; + int dim = tuple->args[0].as()->value; + if (static_cast(dim) >= vinfo.size()) { + vinfo.resize(dim + 1); + } + vinfo[dim].align_factor = tuple->args[1].as()->value; + vinfo[dim].align_offset = tuple->args[2].as()->value; + + return this->VisitStmt(op->body); + } else if (op->attr_key == attr::buffer_bind_scope) { + Array arr = Downcast>(op->node); + ICHECK_EQ(arr.size(), 2U); + Buffer source = Downcast(arr[0]); + Buffer target_with_strides = WithStrides(Downcast(arr[1])); + Buffer source_with_strides = WithStrides(source); + + { + BufferEntry entry; + entry.remap_to = source_with_strides; + entry.in_scope = true; + entry.is_external = false; + buf_map_[source] = entry; + } + + Stmt body = this->VisitStmt(op->body); + + return AttrStmt(Array{source_with_strides, target_with_strides}, op->attr_key, + op->value, body, op->span); + } else { + return StmtExprMutator::VisitStmt_(op); + } + } + + Stmt VisitStmt_(const BufferRealizeNode* op) final { + Buffer key = op->buffer; + Buffer with_strides = WithStrides(op->buffer); + { + BufferEntry entry; + entry.remap_to = with_strides; + entry.in_scope = true; + entry.is_external = false; + buf_map_[key] = entry; + } + + Stmt stmt = StmtExprMutator::VisitStmt_(op); + + buf_map_[key].in_scope = false; + op = stmt.as(); + ICHECK(op); + + return BufferRealize(with_strides, op->bounds, op->condition, op->body, op->span); + } + + PrimExpr VisitExpr_(const BufferLoadNode* op) final { + PrimExpr expr = StmtExprMutator::VisitExpr_(op); + op = expr.as(); + + auto it = buf_map_.find(op->buffer); + ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << op->buffer; + const BufferEntry& e = it->second; + ICHECK(e.in_scope) << "Cannot read a buffer that is already out of scope"; + + return BufferLoad(e.remap_to, op->indices, op->span); + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + + auto it = buf_map_.find(op->buffer); + ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << op->buffer; + const BufferEntry& e = it->second; + ICHECK(e.in_scope) << "Cannot write to a buffer that is already out of scope"; + + return BufferStore(e.remap_to, op->value, op->indices, op->span); + } + + private: + Map updated_extern_buffer_map_; + + struct DimAlignInfo { + int align_factor{0}; + int align_offset{0}; + }; + + // Dimension alignment + std::unordered_map, ObjectPtrHash, ObjectPtrEqual> dim_align_; + + struct BufferEntry { + Buffer remap_to; + bool in_scope; + bool is_external; + }; + + std::unordered_map buf_map_; + + IRVisitorWithAnalyzer* bound_analyzer_; +}; + +/* Use the scope of IterVar to determine storage scope. + * + * For buffers that do not have an explicit storage scope defined, a + * reasonable storage scope may be defined based on the thread scope + * that contains the buffer's allocation. All other buffers without a + * scope are assigned to global scope. + */ +class ThreadScopePropagate : public StmtExprMutator { + public: + explicit ThreadScopePropagate(const Map& extern_buffer_map) { + // External buffers shouldn't be overwritten, even if they have a + // BufferRealizeNode. + for (auto kv : extern_buffer_map) { + external_buffers_.insert(kv.second); + } + } + + PrimExpr VisitExpr_(const VarNode* op) final { + auto it = buf_remap_.find(GetRef(op)); + if (it != buf_remap_.end()) { + return it->second->data; + } else { + return GetRef(op); + } + } + + Stmt VisitStmt_(const AttrStmtNode* op) final { + ICHECK_NE(op->attr_key, attr::buffer_dim_align) + << "StorageFlattener assumes that all buffers have accurate strides, " + << "and all buffer_dim_align annotations are removed. " + << "Please run BufferStrideLegalize first."; + + if (op->attr_key == attr::thread_extent) { + IterVar iv = Downcast(op->node); + ThreadScope ts = ThreadScope::Create(iv->thread_tag); + curr_thread_scope_.push_back(ts); + Stmt stmt = StmtExprMutator::VisitStmt_(op); + curr_thread_scope_.pop_back(); + return stmt; + } else if (op->attr_key == attr::buffer_bind_scope) { + return HandleBufferBindScope(op); + } else { + return StmtExprMutator::VisitStmt_(op); + } + } + + Stmt VisitStmt_(const BufferRealizeNode* op) final { + Var old_var = op->buffer->data; + + // Don't remap buffers that already have an explicit scope, + // or external buffers. + std::string str_scope = GetPtrStorageScope(old_var); + if ((str_scope.length() > 0) || external_buffers_.count(op->buffer)) { + return StmtExprMutator::VisitStmt_(op); + } + + ICHECK_EQ(buf_remap_.count(old_var), 0) + << "Buffer var " << op->buffer->data << " appears in multiple BufferRealize nodes"; + + StorageScope skey; + if (curr_thread_scope_.size() == 0) { + skey.rank = StorageRank::kGlobal; + } else { + skey.rank = runtime::DefaultStorageRank(curr_thread_scope_.back().rank); + } + + auto ptr_type = old_var->type_annotation.as(); + ICHECK(ptr_type); + Var new_var(old_var->name_hint, PointerType(ptr_type->element_type, skey.to_string()), + old_var->span); + + Buffer buf = op->buffer; + buf.CopyOnWrite()->data = new_var; + + buf_remap_[old_var] = buf; + + Stmt body = this->VisitStmt(op->body); + return BufferRealize(buf, op->bounds, op->condition, body, op->span); + } + + PrimExpr VisitExpr_(const BufferLoadNode* op) final { + PrimExpr expr = StmtExprMutator::VisitExpr_(op); + op = expr.as(); + ICHECK(op); + + auto it = buf_remap_.find(op->buffer->data); + if (it != buf_remap_.end()) { + return BufferLoad(it->second, op->indices, op->span); + } else { + return expr; + } + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + ICHECK(op); + + auto it = buf_remap_.find(op->buffer->data); + if (it != buf_remap_.end()) { + return BufferStore(it->second, op->value, op->indices, op->span); + } else { + return stmt; + } + } + + private: + // If the rewritten buffers are part of a buffer_bind_scope, either + // as the buffer view or as the the buffer being viewed, then the + // buffer_bind_scope must be rewritten to refer to the updated + // buffers. + Stmt HandleBufferBindScope(const AttrStmtNode* op) { + Array arr = Downcast>(op->node); + ICHECK_EQ(arr.size(), 2U); + Buffer buffer = Downcast(arr[0]); + ICHECK(buffer.defined()); + Buffer target = Downcast(arr[1]); + ICHECK(target.defined()); + + bool needs_rewrite = false; + + { + auto it = buf_remap_.find(buffer->data); + if (it != buf_remap_.end()) { + needs_rewrite = true; + buffer = it->second; + } + } + + { + auto it = buf_remap_.find(target->data); + if (it != buf_remap_.end()) { + needs_rewrite = true; + target = it->second; + } + } + + if (needs_rewrite) { + Stmt body = this->VisitStmt(op->body); + return AttrStmt(Array{buffer, target}, op->attr_key, op->value, body); + } else { + return StmtExprMutator::VisitStmt_(op); + } + } + + std::unordered_map buf_remap_; + std::unordered_set external_buffers_; + + // The current thread scope. + std::vector curr_thread_scope_; +}; + +/* Map buffer binds to their source buffer + * + * Buffers defined using an attr::buffer_bind_scope annotation are + * views into some linked buffer, potentially into some restricted + * subregion of that buffer. This pass identifies such buffers, then + * rewrites all access of the bound buffers to be access into the + * linked buffer. + */ +class BufferBindUnwrapper : public StmtExprMutator { + public: + explicit BufferBindUnwrapper(const Map& extern_buffer_map, + IRVisitorWithAnalyzer* bound_analyzer) + : bound_analyzer_(bound_analyzer) { + for (auto kv : extern_buffer_map) { + BufferEntry e; + e.buffer = kv.second; + e.external = true; + buf_map_[kv.second.get()] = std::move(e); + } + } + + Stmt VisitStmt_(const StoreNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + auto it = var_remap_.find(op->buffer_var.get()); + if (it != var_remap_.end() && !it->second.same_as(op->buffer_var)) { + // TODO(Lunderberg): Change from warning to error once all mixed + // use of physical/logical layouts is removed. + DLOG(WARNING) << op->buffer_var << " was declared as buffer (buffer_bind_scope), " + << "but is accessed as a pointer (StoreNode)."; + + ICHECK(it->second.as()); + Var new_buf_var = Downcast(it->second); + return Store(new_buf_var, op->value, op->index, op->predicate); + } else { + return stmt; + } + } + + PrimExpr VisitExpr_(const LoadNode* op) final { + PrimExpr expr = StmtExprMutator::VisitExpr_(op); + op = expr.as(); + auto it = var_remap_.find(op->buffer_var.get()); + if (it != var_remap_.end() && !it->second.same_as(op->buffer_var)) { + // TODO(Lunderberg): Change from warning to error once all mixed + // use of physical/logical layouts is removed. + DLOG(WARNING) << op->buffer_var << " was declared as buffer (buffer_bind_scope), " + << "but is accessed as a pointer (LoadNode)."; + + ICHECK(it->second.as()); + Var new_buf_var = Downcast(it->second); + return Load(op->dtype, new_buf_var, op->index, op->predicate); + } else { + return expr; + } + } + + Stmt VisitStmt_(const AttrStmtNode* op) final { + ICHECK_NE(op->attr_key, attr::buffer_dim_align) + << "BufferBindUnwrapper assumes that all buffers have accurate strides, " + << "and all buffer_dim_align annotations are removed. " + << "Please run BufferStrideLegalize first."; + + if (op->attr_key == attr::buffer_bind_scope) { + return HandleBufferBindScope(op); + } else { + return StmtExprMutator::VisitStmt_(op); + } + } + + PrimExpr VisitExpr_(const VarNode* op) final { + auto it = var_remap_.find(op); + if (it != var_remap_.end()) { + return it->second; + } else { + return GetRef(op); + } + } + + Array remap_indices(Array indices, Array begins, + Array extents) { + ICHECK_EQ(begins.size(), extents.size()); + + if (begins.size() == 0) { + return indices; + } + + ICHECK_EQ(begins.size(), indices.size()); + + Array out; + for (size_t i = 0; i < begins.size(); i++) { + out.push_back(begins[i] + indices[i]); + } + return out; + } + + Array remap_bounds(Array bounds, Array begins, Array extents) { + ICHECK_EQ(begins.size(), extents.size()); + + if (begins.size() == 0) { + return bounds; + } + + ICHECK_EQ(begins.size(), bounds.size()); + + Array out; + for (size_t i = 0; i < begins.size(); i++) { + out.push_back(Range::FromMinExtent(bounds[i]->min + begins[i], bounds[i]->extent)); + } + return out; + } + + PrimExpr VisitExpr_(const BufferLoadNode* op) final { + PrimExpr expr = StmtExprMutator::VisitExpr_(op); + op = expr.as(); + + auto it = buf_map_.find(op->buffer.get()); + ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << op->buffer; + const BufferEntry& e = it->second; + ICHECK(e.in_scope) << "Cannot read from buffer " << op->buffer << ", out of scope."; + + if (e.remap) { + return BufferLoad(e.remap->target, + remap_indices(op->indices, e.remap->begins, e.remap->extents), op->span); + } else { + return expr; + } + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + + auto it = buf_map_.find(op->buffer.get()); + ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << op->buffer; + const BufferEntry& e = it->second; + ICHECK(e.in_scope) << "Cannot write to buffer" << op->buffer << ", out of scope."; + + if (e.remap) { + return BufferStore(e.remap->target, op->value, + remap_indices(op->indices, e.remap->begins, e.remap->extents), op->span); + } else { + return stmt; + } + } + + Stmt VisitStmt_(const BufferRealizeNode* op) final { + const auto& key = op->buffer.get(); + + bool is_external = false; + + if (buf_map_.count(key)) { + ICHECK(buf_map_.at(key).external) + << "BufferRealize node for internal buffer " << op->buffer << " occurred multiple times."; + + is_external = true; + } else { + BufferEntry e; + e.bounds = op->bounds; + e.buffer = op->buffer; + buf_map_[key] = std::move(e); + } + + Stmt stmt = StmtExprMutator::VisitStmt_(op); + + if (is_external) { + buf_map_[key].in_scope = false; + } + + return stmt; + } + + Stmt VisitStmt_(const PrefetchNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + op = stmt.as(); + ICHECK(op != nullptr); + + const auto& key = op->buffer.get(); + auto it = buf_map_.find(key); + ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << key; + const BufferEntry& e = it->second; + + ICHECK(e.in_scope) << "Read a buffer that is already out of scope"; + ICHECK_EQ(e.buffer->shape.size(), op->bounds.size()) + << "Prefetch dim should be the same as buffer dim"; + + if (e.remap) { + return Prefetch(e.remap->target, remap_bounds(op->bounds, e.remap->begins, e.remap->extents), + op->span); + } else { + return stmt; + } + } + + private: + // Read the mapping from a buffer view to the actual buffer. This + // allows all later BufferStore/BufferLoad nodes to reference the + // actual buffer, rather than the buffer view. + Stmt HandleBufferBindScope(const AttrStmtNode* op) { + // Unpack information from Attribute node + RemapInfo remap; + + Array arr = Downcast>(op->node); + ICHECK_EQ(arr.size(), 2U); + const Buffer source = Downcast(arr[0]); + ICHECK(source.defined()); + remap.target = Downcast(arr[1]); + ICHECK(remap.target.defined()); + const CallNode* tuple = op->value.as(); + ICHECK(tuple && tuple->op.same_as(builtin::tvm_tuple())); + + for (size_t i = 0; i < tuple->args.size(); i += 2) { + remap.begins.push_back(tuple->args[i]); + remap.extents.push_back(tuple->args[i + 1]); + } + + // Determine bounds in the target buffer + auto it = buf_map_.find(remap.target.get()); + ICHECK(it != buf_map_.end()) << "Cannot define " << source << " as a view into " << remap.target + << ", " << remap.target << " was not defined."; + const BufferEntry& target_info = it->second; + ICHECK(target_info.in_scope) << "Cannot define " << source << " as a view into " << remap.target + << ", " << remap.target << " is out of scope."; + ICHECK_EQ(remap.begins.size(), target_info.buffer->shape.size()) + << "Incorrect number of arguments in buffer_bind_scope attribute. " + << "Expected (min_0, extent_0, min_1, extent_0, ..., min_N, extent_N)."; + + if (target_info.bounds.size() > 0) { + Array mapped_begins; + for (size_t i = 0; i < target_info.buffer->shape.size(); ++i) { + mapped_begins.push_back(remap.begins[i] - target_info.bounds[i]->min); + } + remap.begins = std::move(mapped_begins); + } + + ICHECK(target_info.remap == nullptr) + << "buffer_bind_scope defines " << source << " as a view into " << remap.target + << ", which is itself a buffer view. " + << "Indirect remapping not currently supported."; + + for (size_t i = 0; i < remap.begins.size(); i++) { + remap.begins.Set(i, bound_analyzer_->Simplify(remap.begins[i])); + remap.extents.Set(i, bound_analyzer_->Simplify(remap.extents[i])); + } + + // Add a buffer remap entry + { + BufferEntry source_info; + source_info.buffer = source; + source_info.remap = std::make_unique(remap); + + buf_map_[source.get()] = std::move(source_info); + } + + // Define remappings of any remaining Variables (e.g. Store/Load nodes). + ArgBinder binder(&var_remap_); + + // Define a view that represents the source's view into the target + // buffer. This Buffer object is only used to define the mapping + // to the target buffer, and never actually appears in the TIR + // graph. + Buffer view = remap.target.MakeSlice(remap.begins, remap.extents); + if (source->strides.size() == 0) { + ICHECK_EQ(view->strides.size(), 0U) + << "Cannot bind a compact buffer " << source << " to a strided buffer " << view + << " with strides " << view->strides; + } else { + // Add explicit strides to the view, in order to bind to source.strides[i]. + view = view.MakeStrideView(); + } + + // Bind any variables that reference the view (e.g. elem_offset, + // strides, shape). Pass fuzzy_match=false, because all shape + // transformations should have been handled in + // BufferShapeLegalize. + binder.BindBuffer(source, view, source->name, false); + + // Apply the remaps + Stmt body = op->body; + body = MergeNest(binder.asserts(), body); + body = MergeNest(binder.init_nest(), body); + body = this->VisitStmt(body); + // remove the binds + for (const Var& v : binder.defs()) { + var_remap_.erase(v.get()); + } + return body; + } + + struct RemapInfo { + Buffer target; + Array begins; + Array extents; + }; + + // The buffer entry in the flatten map + struct BufferEntry { + // The storage buffer + Buffer buffer; + // the bounds of realization, can be null, means everything + Region bounds; + // Whether the buffer is external + bool external{false}; + // Whether we are within the allocation scope of the buffer. + bool in_scope{true}; + + // The buffer to which the storage buffer should be remapped. + std::unique_ptr remap{nullptr}; + }; + + // The buffer assignment map + // Variable remap + std::unordered_map var_remap_; + // Buffer map + std::unordered_map buf_map_; + // Analyzer for the variable bounds, used to simplify the bounds populator. We really need the + // analyzer from it. However + IRVisitorWithAnalyzer* bound_analyzer_; +}; + class StorageFlattener : public StmtExprMutator { public: explicit StorageFlattener(const Map& extern_buffer_map, int cache_line_size, @@ -78,6 +1056,15 @@ class StorageFlattener : public StmtExprMutator { } Stmt VisitStmt_(const AttrStmtNode* op) final { + ICHECK_NE(op->attr_key, attr::buffer_dim_align) + << "StorageFlattener assumes that all buffers have accurate strides, " + << "and all buffer_dim_align annotations are removed. " + << "Please run BufferStrideLegalize first."; + + ICHECK_NE(op->attr_key, attr::buffer_bind_scope) + << "StorageFlattener assumes that all buffer binds have already been applied. " + << "Please run BufferBindUnwrapper first."; + if (op->attr_key == attr::double_buffer_scope && op->node->IsInstance()) { auto buffer = Downcast(op->node); Stmt body = this->VisitStmt(op->body); @@ -85,27 +1072,6 @@ class StorageFlattener : public StmtExprMutator { ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << buffer; body = AttrStmt(it->second.buffer->data, op->attr_key, op->value, std::move(body)); return body; - } else if (op->attr_key == attr::thread_extent) { - IterVar iv = Downcast(op->node); - ThreadScope ts = ThreadScope::Create(iv->thread_tag); - curr_thread_scope_.push_back(ts); - Stmt stmt = StmtExprMutator::VisitStmt_(op); - curr_thread_scope_.pop_back(); - return stmt; - } else if (op->attr_key == attr::buffer_bind_scope) { - return HandleBufferBindScope(op); - } else if (op->attr_key == attr::buffer_dim_align) { - auto buffer = Downcast(op->node); - const CallNode* tuple = op->value.as(); - ICHECK(tuple && tuple->op.same_as(builtin::tvm_tuple())); - auto& vinfo = dim_align_[buffer]; - int dim = tuple->args[0].as()->value; - if (static_cast(dim) >= vinfo.size()) { - vinfo.resize(dim + 1); - } - vinfo[dim].align_factor = tuple->args[1].as()->value; - vinfo[dim].align_offset = tuple->args[2].as()->value; - return this->VisitStmt(op->body); } return StmtExprMutator::VisitStmt_(op); } @@ -121,9 +1087,9 @@ class StorageFlattener : public StmtExprMutator { ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << key; const BufferEntry& e = it->second; - ICHECK(!e.released) << "Read a buffer that is already out of scope"; + ICHECK(e.in_scope) << "Cannot write to " << op->buffer << ", out of scope."; - Stmt body = e.buffer.vstore(e.RelIndex(op->indices), op->value); + Stmt body = e.buffer.vstore(op->indices, op->value); if (create_bound_attributes_ && ShapeIsValid(e.buffer->shape)) { shape_collector_.push_back(std::make_pair(e.buffer->data, e.buffer->shape)); } @@ -141,26 +1107,31 @@ class StorageFlattener : public StmtExprMutator { const auto& key = op->buffer; if (buf_map_.count(key)) { - ICHECK(buf_map_.at(key).external); + ICHECK(buf_map_.at(key).external) + << "BufferRealize for internal buffer " << op->buffer << " appears multiple times."; return this->VisitStmt(op->body); } else { // create a buffer entry BufferEntry e; - e.bounds = op->bounds; - Array shape; - for (auto r : e.bounds) { - shape.push_back(r->extent); - } - // deduce current storage scope. - StorageScope skey; - std::string strkey = GetPtrStorageScope(op->buffer->data); - if (strkey.length() == 0) { - if (curr_thread_scope_.size() != 0) { - skey.rank = runtime::DefaultStorageRank(curr_thread_scope_.back().rank); - } - } else { - skey = StorageScope::Create(strkey); + + ICHECK_EQ(op->buffer->shape.size(), op->bounds.size()) + << "Inconsistent buffer shape and realization shape for " << op->buffer; + + for (size_t i = 0; i < op->bounds.size(); i++) { + const auto& bound = op->bounds[i]; + const auto& dim_size = op->buffer->shape[i]; + ICHECK(is_zero(bound_analyzer_->Simplify(bound->min))) + << "Buffer " << op->buffer << " has realization bounds that do not start at zero. " + << "Please run BufferShapeLegalize first."; + ICHECK(is_one(bound_analyzer_->Simplify(bound->extent == dim_size))) + << "Buffer " << op->buffer + << " has realization extent that does not match its size. " + "Please run BufferShapeLegalize first."; } + + Array shape = op->buffer->shape; + StorageScope skey = StorageScope::Create(GetPtrStorageScope(op->buffer->data)); + // use small alignment for small arrays auto dtype = op->buffer->dtype; int32_t const_size = AllocateNode::constant_allocation_size(shape); @@ -173,36 +1144,14 @@ class StorageFlattener : public StmtExprMutator { << "Allocation exceed bound of memory tag " << skey.to_string(); } } - Array strides; - if (dim_align_.count(key) != 0 && shape.size() != 0) { - std::vector rstrides; - const std::vector& avec = dim_align_[key]; - int first_dim = 0; - PrimExpr stride = make_const(shape[first_dim].dtype(), 1); - for (size_t i = shape.size(); i != 0; --i) { - size_t dim = i - 1; - if (dim < avec.size() && avec[dim].align_factor != 0) { - PrimExpr factor = make_const(stride.dtype(), avec[dim].align_factor); - PrimExpr offset = make_const(stride.dtype(), avec[dim].align_offset); - stride = stride + indexmod(factor + offset - indexmod(stride, factor), factor); - stride = bound_analyzer_->Simplify(stride); - } - rstrides.push_back(stride); - stride = stride * shape[dim]; - } - strides = Array(rstrides.rbegin(), rstrides.rend()); - } + Array strides = op->buffer->strides; - auto* ptr_type = op->buffer->data->type_annotation.as(); - ICHECK(ptr_type); - auto new_var = - Var(op->buffer->data->name_hint, PointerType(ptr_type->element_type, skey.to_string())); - e.buffer = Buffer(new_var, op->buffer->dtype, shape, strides, PrimExpr(), op->buffer->name, - align, 0, kDefault); + e.buffer = Buffer(op->buffer->data, op->buffer->dtype, shape, strides, PrimExpr(), + op->buffer->name, align, 0, kDefault); buf_map_[key] = e; Stmt body = this->VisitStmt(op->body); - buf_map_[key].released = true; + buf_map_[key].in_scope = false; Stmt ret; DataType storage_type = e.buffer->dtype; @@ -264,12 +1213,12 @@ class StorageFlattener : public StmtExprMutator { auto it = buf_map_.find(key); ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << key; const BufferEntry& e = it->second; - ICHECK(!e.released) << "Read a buffer that is already out of scope"; + ICHECK(e.in_scope) << "Cannot read to " << op->buffer << ", out of scope."; if (create_bound_attributes_ && ShapeIsValid(e.buffer->shape)) { shape_collector_.push_back(std::make_pair(e.buffer->data, e.buffer->shape)); } - return e.buffer.vload(e.RelIndex(op->indices), e.buffer->dtype); + return e.buffer.vload(op->indices, e.buffer->dtype); } Stmt VisitStmt_(const PrefetchNode* op) final { @@ -282,7 +1231,7 @@ class StorageFlattener : public StmtExprMutator { ICHECK(it != buf_map_.end()) << "Cannot find allocated buffer for " << key; const BufferEntry& e = it->second; - ICHECK(!e.released) << "Read a buffer that is already out of scope"; + ICHECK(e.in_scope) << "Cannot prefetch " << op->buffer << ", out of scope."; ICHECK_EQ(e.buffer->shape.size(), op->bounds.size()) << "Prefetch dim should be the same as buffer dim"; @@ -315,7 +1264,7 @@ class StorageFlattener : public StmtExprMutator { if (i < starts) { stmt = For(vars[i], 0, op->bounds[i]->extent, ForKind::kSerial, stmt); } else { - PrimExpr load = e.buffer.vload(e.RelIndex(args), e.buffer->dtype); + PrimExpr load = e.buffer.vload(args, e.buffer->dtype); PrimExpr address = Call(DataType::Handle(), builtin::address_of(), {load}); PrimExpr prefetch = Call(op->buffer->dtype, builtin::prefetch(), {address, 0, 3, 1}); stmt = Evaluate(prefetch); @@ -327,107 +1276,24 @@ class StorageFlattener : public StmtExprMutator { } PrimExpr VisitExpr_(const ProducerLoadNode* op) final { - LOG(FATAL) << "ProducerLoad cannot appear in a valid TIR PrimFunc."; + LOG(FATAL) << "ProducerLoad cannot appear in a valid TIR PrimFunc. " + << "Please run SchedulePostProcToPrimFunc first."; return PrimExpr(); } Stmt VisitStmt_(const ProducerStoreNode* op) final { - LOG(FATAL) << "Cannot handle Provide " - << " please run SchedulePostProcToPrimFunc first"; + LOG(FATAL) << "ProducerStore cannot appear in a valid TIR PrimFunc. " + << "Please run SchedulePostProcToPrimFunc first."; return Stmt(); } Stmt VisitStmt_(const ProducerRealizeNode* op) final { - LOG(FATAL) << "Cannot handle Realize " - << " please run SchedulePostProcToPrimFunc first"; + LOG(FATAL) << "ProducerRealize cannot appear in a valid TIR PrimFunc. " + << "Please run SchedulePostProcToPrimFunc first."; return Stmt(); } private: - // The specific tensor data layout is not determined before - // StorageFlatten pass. We use buffer_bind_scope - // to specify before hand we want to bind a subregion - // of tensor to a symbolic buffer, which get used in extern. - // - // Example: - // - // realize A in range [i*4, extent=10) { - // bind Ab to A in [i*4+1, extent=4) { - // call_func(Ab.ptr, Ab.shape[0]) - // } - // } - // - // After StorageFlatten - // - // alloc A[10] - // call(A + 1, 4) - // - // Buffer is a protocol to declare specific - // data layout and shape we expect. - // So this function need to check: - // - If the bind range is within the realize range - // - If we can match the requirement of buffer - // - Remap variables such as Ab.ptr to the actual value. - // - // Here are a few possible failure cases: - // - Buffer is declared to have constant shape, - // but we try to bind it to a different one. - // - Buffer is declared to be compact(no strides) - // but this binded region is a subregion of - // a matrix(tensor), which means it requires strides. - // - // We do support a few relaxed case, such as bindingx - // region with shape [1, 1, n, m] to buffer with shape [n, m] - Stmt HandleBufferBindScope(const AttrStmtNode* op) { - Array arr = Downcast>(op->node); - ICHECK_EQ(arr.size(), 2U); - const BufferNode* buffer = arr[0].as(); - const BufferNode* target = arr[1].as(); - const CallNode* tuple = op->value.as(); - ICHECK(buffer && target); - ICHECK(tuple && tuple->op.same_as(builtin::tvm_tuple())); - auto key = GetRef(target); - - auto it = buf_map_.find(key); - ICHECK(it != buf_map_.end()) << "Cannot find buffer of " << key; - const BufferEntry& be = it->second; - ICHECK(!be.released); - ICHECK_EQ(tuple->args.size(), be.buffer->shape.size() * 2); - Array begins, extents; - if (be.bounds.size() != 0) { - ICHECK_EQ(tuple->args.size(), be.bounds.size() * 2); - for (size_t i = 0; i < be.buffer->shape.size(); ++i) { - begins.push_back(tuple->args[2 * i] - be.bounds[i]->min); - extents.push_back(tuple->args[2 * i + 1]); - } - } else { - for (size_t i = 0; i < tuple->args.size(); i += 2) { - begins.push_back(tuple->args[i]); - auto new_extent = bound_analyzer_->Simplify(tuple->args[i + 1]); - extents.push_back(new_extent); - } - } - Buffer slice = be.buffer.MakeSlice(begins, extents); - if (buffer->strides.size() == 0) { - ICHECK_EQ(slice->strides.size(), 0U) - << "Trying to bind compact buffer to strided one strides=" << slice->strides; - } else { - slice = slice.MakeStrideView(); - } - // start binding - ArgBinder binder(&var_remap_); - binder.BindBuffer(Downcast(arr[0]), slice, buffer->name, true); - // Apply the remaps - Stmt body = MergeNest(binder.asserts(), op->body); - body = MergeNest(binder.init_nest(), body); - body = this->VisitStmt(body); - // remove the binds - for (const Var& v : binder.defs()) { - var_remap_.erase(v.get()); - } - return body; - } - // The buffer entry in the flatten map struct DimAlignInfo { int align_factor{0}; @@ -437,25 +1303,10 @@ class StorageFlattener : public StmtExprMutator { struct BufferEntry { // the buffer of storage Buffer buffer; - // the bounds of realization, can be null, means everything - Region bounds; // Whether the buffer is external bool external{false}; - // Whether we are out of allocation bounds and buffer get released. - bool released{false}; - // relative index - inline Array RelIndex(Array args) const { - if (bounds.size() != 0) { - Array index; - ICHECK_EQ(bounds.size(), args.size()); - for (size_t i = 0; i < bounds.size(); ++i) { - index.push_back(args[i] - bounds[i]->min); - } - return index; - } else { - return args; - } - } + // Whether the buffer is currently in scope. + bool in_scope{true}; }; bool ShapeIsValid(const Array& shape) { @@ -484,10 +1335,6 @@ class StorageFlattener : public StmtExprMutator { std::unordered_map var_remap_; // Buffer map std::unordered_map buf_map_; - // Dimension alignment - std::unordered_map, ObjectPtrHash, ObjectPtrEqual> dim_align_; - // The current thread scope. - std::vector curr_thread_scope_; // Collects shapes. std::vector>> shape_collector_; // bounds populator. We really need the analyzer from it. @@ -499,6 +1346,68 @@ class StorageFlattener : public StmtExprMutator { bool create_bound_attributes_{false}; }; +/*! + * \brief Simplify assert statements. + * + * If an assert statement can be statically verified to be true, + * remove the assert statement. Otherwise, keep the assert statement + * unmodified. + */ +class AssertSimplifier : public StmtMutator { + public: + explicit AssertSimplifier(IRVisitorWithAnalyzer* bound_analyzer) + : bound_analyzer_(bound_analyzer) {} + + Stmt VisitStmt_(const AssertStmtNode* op) final { + Stmt stmt = StmtMutator::VisitStmt_(op); + op = stmt.as(); + + PrimExpr condition = bound_analyzer_->Simplify(op->condition); + if (is_one(condition)) { + return op->body; + } + + return stmt; + } + + private: + IRVisitorWithAnalyzer* bound_analyzer_; +}; + +// The specific tensor data layout is not determined before +// StorageFlatten pass. We use buffer_bind_scope +// to specify before hand we want to bind a subregion +// of tensor to a symbolic buffer, which get used in extern. +// +// Example: +// +// realize A in range [i*4, extent=10) { +// bind Ab to A in [i*4+1, extent=4) { +// call_func(Ab.ptr, Ab.shape[0]) +// } +// } +// +// After StorageFlatten +// +// alloc A[10] +// call(A + 1, 4) +// +// Buffer is a protocol to declare specific +// data layout and shape we expect. +// So this function need to check: +// - If the bind range is within the realize range +// - If we can match the requirement of buffer +// - Remap variables such as Ab.ptr to the actual value. +// +// Here are a few possible failure cases: +// - Buffer is declared to have constant shape, +// but we try to bind it to a different one. +// - Buffer is declared to be compact(no strides) +// but this binded region is a subregion of +// a matrix(tensor), which means it requires strides. +// +// We do support a few relaxed case, such as binding a +// region with shape [1, 1, n, m] to buffer with shape [n, m] PrimFunc StorageFlatten(PrimFunc func, int cache_line_size, bool create_bound_attributes) { // Only apply this pass to TIR from TE schedules Optional from_legacy_te_schedule = func->GetAttr("from_legacy_te_schedule", Bool(false)); @@ -507,8 +1416,22 @@ PrimFunc StorageFlatten(PrimFunc func, int cache_line_size, bool create_bound_at IRVisitorWithAnalyzer bound_analyzer; bound_analyzer(fptr->body); + + fptr->body = BufferShapeLegalize(fptr->buffer_map, &bound_analyzer)(std::move(fptr->body)); + + auto stride_legalize = BufferStrideLegalize(fptr->buffer_map, &bound_analyzer); + fptr->body = stride_legalize(std::move(fptr->body)); + fptr->buffer_map = stride_legalize.UpdatedExternBufferMap(); + + fptr->body = ThreadScopePropagate(fptr->buffer_map)(std::move(fptr->body)); + + fptr->body = BufferBindUnwrapper(fptr->buffer_map, &bound_analyzer)(std::move(fptr->body)); + fptr->body = StorageFlattener(fptr->buffer_map, cache_line_size, create_bound_attributes, &bound_analyzer)(std::move(fptr->body)); + + fptr->body = AssertSimplifier(&bound_analyzer)(std::move(fptr->body)); + return func; } else { return func;