vinx13 commented on a change in pull request #9680:
URL: https://github.com/apache/tvm/pull/9680#discussion_r766036466
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -471,6 +480,62 @@ Doc TVMScriptPrinter::PrintMatchBufferRegion(const
MatchBufferRegionNode* op) {
return doc;
}
+// check if all arguments, except the first two, are specified for
T.match_buffer
+// if not, then this match buffer is printed out as T.buffer in prim_func
arguments
+bool TVMScriptPrinter::IsSimpleBuffer(const Buffer& buf) {
+ if (memo_var_.find(buf->data) != memo_var_.end()) {
+ return false;
+ }
+ if (!buf->strides.empty()) {
+ return false;
+ }
+ if (buf->elem_offset->IsInstance<VarNode>()) {
+ Var elem_offset = Downcast<Var>(buf->elem_offset);
+ if (memo_var_.find(elem_offset) != memo_var_.end()) {
+ return false;
+ }
+ } else if (buf->elem_offset->IsInstance<IntImmNode>()) {
+ IntImm elem_offset = Downcast<IntImm>(buf->elem_offset);
+ if (elem_offset->value != 0) {
+ return false;
+ }
+ }
+ if (buf.scope() != "global") {
+ return false;
+ }
+ if (buf->data_alignment != runtime::kAllocAlignment) {
+ return false;
+ }
+ if (buf->offset_factor != 1) {
+ return false;
+ }
+ if (buf->buffer_type != BufferType::kDefault) {
+ return false;
+ }
+ return true;
+}
+
+Doc TVMScriptPrinter::MatchBufferDeclaration(const Buffer& buffer) {
+ Doc doc;
+ doc << tir_prefix_ << ".Buffer[" <<
PrintTuple(buffer->shape.as<ArrayNode>());
+ doc << ", " << PrintDType(buffer->dtype) << "]";
+ return doc;
+}
+
+// print array out as tuple with parentheses
+Doc TVMScriptPrinter::PrintTuple(const ArrayNode* op) {
+ Doc doc;
+ doc << '(';
+ for (size_t i = 0; i < op->size(); ++i) {
+ if (i != 0) {
+ doc << ", ";
+ }
+ doc << Print(op->at(i));
+ }
+ doc << ')';
Review comment:
if tuple size is 1, extra `,` is needed at the end, i.e. `(n,)` instead
of `(n)`
##########
File path: tests/python/unittest/test_tvmscript_roundtrip.py
##########
@@ -79,6 +79,7 @@ def mmult(A: T.handle, B: T.handle, C: T.handle) -> None:
def test_opt_gemm_normalize():
mod = Module1
+ print(mod.script(show_meta=True))
Review comment:
remove this
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -220,6 +221,14 @@ class TVMScriptPrinter : public StmtFunctor<Doc(const
Stmt&)>,
Doc AllocBuf(const Buffer& buffer);
void TryDeallocVar(const Var& var);
bool ContainsOptionalInfo(const Stmt& stmt);
+ /*! Helper function for match buffer printing */
+ /*!
+ * @brief check if a T.match_buffer decl is syntax sugarred
Review comment:
uses `\brief` for consistency, the comment should also be made more clear
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -1220,6 +1300,19 @@ Doc TVMScriptPrinter::PrintPrimFunc(const PrimFunc&
primFunc) {
std::vector<Doc> params;
for (const auto& param : op->params) {
var_not_in_headers_.insert(param.get());
+ auto it = op->buffer_map.find(param);
+ // check if this param is a T.handle
+ if (it != op->buffer_map.end()) {
+ // check if this match_buffer has only the first two arguments specified
+ const auto buf = (*it).second;
+ if (IsSimpleBuffer(buf)) {
+ buf_not_in_headers_.insert(buf.get());
+ Doc buf_param_doc;
+ buf_param_doc << buf->name << ": " << MatchBufferDeclaration(buf);
Review comment:
This is problematic. See the impl of `AllocBuf` and
`AllocBufferDeclaration`, I can see at least
`memo_buf_`, `memo_var_` are not correctly updated if the buffer doesn't go
through these two functions. Possible solution is to unify the new behavior
into these two functions
--
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]
