vinx13 commented on a change in pull request #9680:
URL: https://github.com/apache/tvm/pull/9680#discussion_r766291329
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -1218,8 +1296,21 @@ Doc TVMScriptPrinter::PrintPrimFunc(const PrimFunc&
primFunc) {
doc << "def " << (func2var_.find(op) == func2var_.end() ? "func" :
func2var_[op]->name_hint)
<< "(";
std::vector<Doc> params;
+ std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual> simple_buf;
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;
Review comment:
```suggestion
const Buffer& buf = (*it).second;
```
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -220,8 +221,13 @@ class TVMScriptPrinter : public StmtFunctor<Doc(const
Stmt&)>,
Doc AllocBuf(const Buffer& buffer);
void TryDeallocVar(const Var& var);
bool ContainsOptionalInfo(const Stmt& stmt);
-
- /*! Helper functions for loop printing. */
+ /*!
+ * @brief check if a T.match_buffer decl has only first two arguments
specified
Review comment:
```suggestion
* \brief check if a buffer declaration has only 'shape' and 'dtype'
arguments specified
```
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -220,8 +221,13 @@ class TVMScriptPrinter : public StmtFunctor<Doc(const
Stmt&)>,
Doc AllocBuf(const Buffer& buffer);
void TryDeallocVar(const Var& var);
bool ContainsOptionalInfo(const Stmt& stmt);
-
- /*! Helper functions for loop printing. */
+ /*!
+ * @brief check if a T.match_buffer decl has only first two arguments
specified
+ * \param buffer The match buffer to be checked
+ */
+ bool IsSimpleBuffer(const Buffer& buffer);
+ Doc MatchBufferDeclaration(const Buffer& buffer);
Review comment:
```suggestion
Doc PrintInlineBufferBind(const Buffer& buffer);
```
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -220,8 +221,13 @@ class TVMScriptPrinter : public StmtFunctor<Doc(const
Stmt&)>,
Doc AllocBuf(const Buffer& buffer);
void TryDeallocVar(const Var& var);
bool ContainsOptionalInfo(const Stmt& stmt);
-
- /*! Helper functions for loop printing. */
Review comment:
keep this
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -471,6 +477,63 @@ 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;
+ }
Review comment:
```suggestion
return false;
```
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -1229,9 +1320,12 @@ Doc TVMScriptPrinter::PrintPrimFunc(const PrimFunc&
primFunc) {
for (const auto& param : op->params) {
auto it = op->buffer_map.find(param);
if (it == op->buffer_map.end()) continue;
- buf_not_in_headers_.insert((*it).second.get());
- body << Print((*it).second) << " = " << tir_prefix_ << ".match_buffer(";
- body << Print((*it).first) << ", " << memo_buf_decl_[(*it).second];
+ auto buf = (*it).second;
+ auto search = simple_buf.find(buf);
+ if (search != simple_buf.end()) continue;
Review comment:
```suggestion
if (simple_buf.count(buf)) continue;
```
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -1229,9 +1320,12 @@ Doc TVMScriptPrinter::PrintPrimFunc(const PrimFunc&
primFunc) {
for (const auto& param : op->params) {
auto it = op->buffer_map.find(param);
if (it == op->buffer_map.end()) continue;
- buf_not_in_headers_.insert((*it).second.get());
- body << Print((*it).second) << " = " << tir_prefix_ << ".match_buffer(";
- body << Print((*it).first) << ", " << memo_buf_decl_[(*it).second];
+ auto buf = (*it).second;
Review comment:
avoid copying
```suggestion
const Buffer& buf = (*it).second;
```
--
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]
