Hzfengsy commented on a change in pull request #7630:
URL: https://github.com/apache/tvm/pull/7630#discussion_r593894641
##########
File path: src/printer/tvmscript_printer.cc
##########
@@ -713,6 +782,88 @@ Doc TVMScriptPrinter::VisitStmt_(const BufferStoreNode*
op) {
return doc;
}
+Doc TVMScriptPrinter::VisitStmt_(const BlockRealizeNode* op) {
+ const auto* block_op = op->block.as<BlockNode>();
+ // print block name and block vars
+ Doc doc;
+ doc << "with tir.block([";
+ std::vector<Doc> block_var_docs;
+ for (const auto& iter_var : block_op->iter_vars) {
+ Doc block_var_doc;
+ if (is_zero(iter_var->dom->min) && iter_var->iter_type == kDataPar) {
+ block_var_doc << Print(iter_var->dom->extent);
+ } else {
+ block_var_doc << "tir.";
+ switch (iter_var->iter_type) {
+ case kDataPar:
+ block_var_doc << "range";
+ break;
+ case kCommReduce:
+ block_var_doc << "reduce_axis";
+ break;
+ case kOrdered:
+ block_var_doc << "scan_axis";
+ break;
+ case kOpaque:
+ block_var_doc << "opaque_axis";
+ break;
+ default:
+ LOG(FATAL) << "Unknown block var iter type";
Review comment:
It is just an enum class and only able to print the number (such as 0,
1) rather than its name. On the other side, I don't think we need such printing
because it is rare to trigger this error.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
