| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| */ |
| #include "./utils.h" |
| |
| namespace tvm { |
| namespace script { |
| namespace printer { |
| |
| Doc PrintBlock(IRDocsifier d, tirx::SBlock block, AccessPath block_p, // |
| ffi::Optional<tirx::SBlockRealize> opt_realize, |
| ffi::Optional<AccessPath> opt_realize_p) { |
| With<TIRFrame> frame(d, block); |
| TVM_FFI_ICHECK_EQ(opt_realize.has_value(), opt_realize_p.has_value()); |
| const tirx::SBlockRealizeNode* realize = |
| opt_realize.has_value() ? opt_realize.value().get() : nullptr; |
| AccessPath realize_p = *opt_realize_p; |
| |
| // Step 1. Handle block var and block bindings |
| // Step 1.1. Obtain all loop var defined along path |
| std::unordered_map<const tirx::VarNode*, tirx::For> loop_vars; |
| for (Frame f : d->frames) { |
| if (const auto* tir_f = f.as<TIRFrameNode>()) { |
| if (auto for_loop = tir_f->tirx.as<tirx::For>()) { |
| for (ffi::Optional<tirx::For> loop = for_loop; loop; |
| loop = loop.value()->body.as<tirx::For>()) { |
| loop_vars.insert(std::make_pair(loop.value()->loop_var.get(), loop.value())); |
| } |
| } |
| } |
| } |
| |
| std::vector<int> remap_vars_indices; |
| auto add_remapped_iter_var = [&](int i) -> bool { |
| if (realize && d->cfg->syntax_sugar) { |
| tirx::ExprDeepEqual expr_equal; |
| tirx::IterVar iter_var = block->iter_vars[i]; |
| PrimExpr value = realize->iter_values[i]; |
| if (iter_var->iter_type == tirx::IterVarType::kDataPar || |
| iter_var->iter_type == tirx::IterVarType::kCommReduce) { |
| if (auto var = value.as<tirx::PrimVar>()) { |
| if (loop_vars.count(var.value().get())) { |
| tirx::For for_loop = loop_vars.at(var.value().get()); |
| if (expr_equal(for_loop->min, iter_var->dom->min) && |
| expr_equal(for_loop->extent, iter_var->dom->extent)) { |
| remap_vars_indices.push_back(i); |
| return true; |
| } |
| } |
| } |
| } |
| } |
| return false; |
| }; |
| |
| auto print_single_iter_var = [&](int i) { |
| tirx::IterVar iter_var = block->iter_vars[i]; |
| AccessPath iter_var_p = block_p->Attr("iter_var")->ArrayItem(i); |
| ExprDoc rhs = TIR(d, "axis"); |
| if (iter_var->iter_type == tirx::IterVarType::kDataPar) { |
| rhs = rhs->Attr("spatial"); |
| } else if (iter_var->iter_type == tirx::IterVarType::kCommReduce) { |
| rhs = rhs->Attr("reduce"); |
| } else if (iter_var->iter_type == tirx::IterVarType::kOrdered) { |
| rhs = rhs->Attr("scan"); |
| } else if (iter_var->iter_type == tirx::IterVarType::kOpaque) { |
| rhs = rhs->Attr("opaque"); |
| } else { |
| TVM_FFI_THROW(ValueError) << "Unknown IterVarType in block signature: " |
| << tirx::IterVarType2String(iter_var->iter_type); |
| } |
| ExprDoc dom{ffi::UnsafeInit()}; |
| if (tirx::is_zero(iter_var->dom->min)) { |
| ExprDoc extent = d->AsDoc<ExprDoc>(iter_var->dom->extent, // |
| iter_var_p->Attr("dom")->Attr("extent")); |
| dom = extent; |
| } else { |
| ExprDoc min = d->AsDoc<ExprDoc>(iter_var->dom->min, iter_var_p->Attr("dom")->Attr("min")); |
| ExprDoc max = d->AsDoc<ExprDoc>(iter_var->dom->min + iter_var->dom->extent, |
| iter_var_p->Attr("dom")->Attr("extent")); |
| dom = TupleDoc({min, max}); |
| } |
| if (realize) { |
| ExprDoc binding = d->AsDoc<ExprDoc>(realize->iter_values[i], // |
| realize_p->Attr("iter_values")->ArrayItem(i)); |
| rhs = rhs->Call({dom, binding}); |
| } else { |
| rhs = rhs->Call({dom}); |
| } |
| (*frame)->stmts.push_back(AssignDoc(DefineVar(iter_var->var, *frame, d), rhs, std::nullopt)); |
| }; |
| |
| auto print_remapped_iter_var = [&]() { |
| if (remap_vars_indices.size()) { |
| int m = remap_vars_indices.size(); |
| if (m == 1) { |
| print_single_iter_var(remap_vars_indices[0]); |
| remap_vars_indices.clear(); |
| return; |
| } |
| ffi::Array<ExprDoc> lhs; |
| ffi::Array<ExprDoc> loop_var_doc; |
| lhs.reserve(m); |
| loop_var_doc.reserve(m); |
| std::string binding_type = ""; |
| ffi::Array<AccessPath> binding_paths; |
| for (int i : remap_vars_indices) { |
| tirx::IterVar iter_var = block->iter_vars[i]; |
| AccessPath iter_var_p = block_p->Attr("iter_vars")->ArrayItem(i); |
| lhs.push_back(DefineVar(iter_var->var, *frame, d)); |
| loop_var_doc.push_back(d->AsDoc<ExprDoc>(realize->iter_values[i], |
| realize_p->Attr("iter_values")->ArrayItem(i))); |
| binding_paths.push_back(iter_var_p->Attr("iter_type")); |
| binding_type += iter_var->iter_type == tirx::IterVarType::kDataPar ? "S" : "R"; |
| } |
| ExprDoc rhs = TIR(d, "axis")->Attr("remap"); |
| ExprDoc binding_str = LiteralDoc::Str(binding_type, std::nullopt); |
| binding_str->source_paths = std::move(binding_paths); |
| rhs = rhs->Call({binding_str, ListDoc(loop_var_doc)}); |
| (*frame)->stmts.push_back(AssignDoc(TupleDoc(lhs), rhs, std::nullopt)); |
| remap_vars_indices.clear(); |
| } |
| }; |
| |
| // Step 1.2. Construct all block var bindings |
| int n_vars = block->iter_vars.size(); |
| for (int i = 0; i < n_vars; ++i) { |
| if (!add_remapped_iter_var(i)) { |
| print_remapped_iter_var(); |
| print_single_iter_var(i); |
| } |
| } |
| print_remapped_iter_var(); |
| |
| // Step 2. Handle block predicate |
| if (realize) { |
| PrimType predicate_ty = realize->predicate.ty(); |
| TVM_FFI_ICHECK(realize->predicate.defined() && |
| predicate_ty.MatchesCode(DLDataTypeCode::kDLBool)); |
| if (!tirx::is_one(realize->predicate)) { |
| (*frame)->stmts.push_back(ExprStmtDoc( |
| TIR(d, "where") |
| ->Call({d->AsDoc<ExprDoc>(realize->predicate, realize_p->Attr("predicate"))}))); |
| } |
| } |
| // Step 3. Handle block read/write regions |
| { |
| ffi::Array<ExprDoc> reads; |
| for (int i = 0, n = block->reads.size(); i < n; ++i) { |
| reads.push_back(d->AsDoc<ExprDoc>(block->reads[i], block_p->Attr("reads")->ArrayItem(i))); |
| } |
| (*frame)->stmts.push_back(ExprStmtDoc(TIR(d, "reads")->Call(reads))); |
| ffi::Array<ExprDoc> writes; |
| for (int i = 0, n = block->writes.size(); i < n; ++i) { |
| writes.push_back(d->AsDoc<ExprDoc>(block->writes[i], block_p->Attr("writes")->ArrayItem(i))); |
| } |
| (*frame)->stmts.push_back(ExprStmtDoc(TIR(d, "writes")->Call(writes))); |
| } |
| // Step 4. Handle block attributes |
| if (!block->annotations.empty()) { |
| (*frame)->stmts.push_back(ExprStmtDoc( |
| TIR(d, "sblock_attr") |
| ->Call({d->AsDoc<ExprDoc>(block->annotations, block_p->Attr("annotations"))}))); |
| } |
| // Step 5. Handle `alloc_buffer` |
| for (int i = 0, n = block->alloc_buffers.size(); i < n; ++i) { |
| tirx::BufferVar buffer = block->alloc_buffers[i]; |
| AccessPath buffer_p = block_p->Attr("alloc_buffers")->ArrayItem(i); |
| IdDoc lhs = DefineBuffer(buffer, *frame, d); |
| ExprDoc rhs = BufferDecl(buffer, "sblock_alloc_buffer", {}, buffer_p, *frame, d, |
| BufferVarDefinition::DataPointer); |
| (*frame)->stmts.push_back(AssignDoc(lhs, rhs, std::nullopt)); |
| } |
| // Step 6. Handle `match_buffer` |
| for (int i = 0, n = block->match_buffers.size(); i < n; ++i) { |
| tirx::MatchBufferRegion buffer_region = block->match_buffers[i]; |
| AccessPath buffer_region_p = block_p->Attr("match_buffers")->ArrayItem(i); |
| StmtDoc doc = d->AsDoc<StmtDoc>(buffer_region, buffer_region_p); |
| (*frame)->stmts.push_back(doc); |
| } |
| // Step 7. Handle init block |
| if (block->init.has_value()) { |
| tirx::Stmt init = block->init.value(); |
| With<TIRFrame> init_frame(d, init); |
| AsDocBody(init, block_p->Attr("init"), init_frame->get(), d); |
| (*frame)->stmts.push_back( |
| ScopeDoc(std::nullopt, TIR(d, "init")->Call({}), (*init_frame)->stmts)); |
| } |
| // Step 8. Handle block body |
| AsDocBody(block->body, block_p->Attr("body"), frame->get(), d); |
| ffi::Array<ffi::String> kwargs_keys; |
| ffi::Array<ExprDoc> kwargs_values; |
| if (!realize) { |
| kwargs_keys.push_back("no_realize"); |
| kwargs_values.push_back(LiteralDoc::Boolean(true, std::nullopt)); |
| } |
| return ScopeDoc(std::nullopt, |
| TIR(d, "sblock") // |
| ->Call({LiteralDoc::Str(block->name_hint, block_p->Attr("name_hint"))}, |
| kwargs_keys, kwargs_values), |
| (*frame)->stmts); |
| } |
| |
| TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
| .set_dispatch<tirx::SBlockRealize>( |
| "", [](tirx::SBlockRealize realize, AccessPath p, IRDocsifier d) -> Doc { |
| Doc doc = PrintBlock(d, realize->block, p->Attr("block"), realize, p); |
| // since we do not have d->AsDoc for realize->block, |
| // we should add possible doc decoration manually. |
| AddDocDecoration<ScopeDoc>(doc, realize->block, p->Attr("block"), d->cfg); |
| return doc; |
| }); |
| |
| TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
| .set_dispatch<tirx::SBlock>("", [](tirx::SBlock block, AccessPath p, IRDocsifier d) -> Doc { |
| return PrintBlock(d, block, p, std::nullopt, std::nullopt); |
| }); |
| |
| TVM_REGISTER_SCRIPT_AS_REPR(tirx::SBlockNode, ReprPrintTIR); |
| TVM_REGISTER_SCRIPT_AS_REPR(tirx::SBlockRealizeNode, ReprPrintTIR); |
| |
| TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
| .set_dispatch<tirx::ScopeIdDefStmt>( |
| "", [](tirx::ScopeIdDefStmt stmt, AccessPath p, IRDocsifier d) -> Doc { |
| // Render as ``(var1, var2, ...) = T.cta_id([ext], preferred=[...])`` |
| // (or the appropriate API name for the binding). |
| TVM_FFI_ICHECK(!d->frames.empty()); |
| tirx::ScopeIdDef def = stmt->def; |
| AccessPath def_p = p->Attr("def"); |
| ffi::Array<ExprDoc> lhs; |
| for (auto scope_id : def->def_ids) { |
| lhs.push_back(DefineVar(scope_id, d->frames.back(), d)); |
| } |
| ffi::Array<ExprDoc> rhs_args; |
| if (def->scope != tirx::ScopeBinding::kClusterCtaPair && def->extents.has_value()) { |
| rhs_args.push_back(d->AsDoc<ExprDoc>(def->extents.value(), def_p->Attr("extents"))); |
| } |
| ffi::Array<ffi::String> kwarg_keys; |
| ffi::Array<ExprDoc> kwarg_vals; |
| if (def->preferred_extents.has_value()) { |
| kwarg_keys.push_back("preferred"); |
| kwarg_vals.push_back(d->AsDoc<ExprDoc>(def->preferred_extents.value(), |
| def_p->Attr("preferred_extents"))); |
| } |
| // The scope-id dtype is independent of the extents, so it has to be printed |
| // explicitly whenever it is not the int32 default in order to round-trip. |
| if (!def->def_ids.empty()) { |
| PrimType scope_id_ty = def->def_ids[0].ty(); |
| for (auto scope_id : def->def_ids) { |
| TVM_FFI_ICHECK(scope_id.ty() == scope_id_ty) |
| << "mixed scope-id dtypes are unsupported, got " << scope_id_ty << " and " |
| << scope_id.ty(); |
| } |
| if (scope_id_ty != PrimType::Int(32)) { |
| kwarg_keys.push_back("dtype"); |
| kwarg_vals.push_back( |
| LiteralDoc::Str(DType2Str(scope_id_ty->dtype), def_p->Attr("def_ids"))); |
| } |
| } |
| ExprDoc rhs = TIR(d, ScopeIdApiName(def->scope))->Call(rhs_args, kwarg_keys, kwarg_vals); |
| return AssignDoc(TupleDoc(lhs), rhs, std::nullopt); |
| }); |
| |
| TVM_SCRIPT_REPR(tirx::ScopeIdDefStmtNode, ReprPrintTIR); |
| |
| TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
| .set_dispatch<tirx::ExecScope>( |
| "", [](tirx::ExecScope exec_scope, AccessPath p, IRDocsifier d) -> Doc { |
| Doc doc = |
| TIR(d, "ExecScope")->Call({LiteralDoc::Str(exec_scope->name(), p->Attr("name"))}); |
| return doc; |
| }); |
| TVM_SCRIPT_REPR(tirx::ExecScopeNode, ReprPrintTIR); |
| |
| TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
| .set_dispatch<tirx::ScopeIdDef>( |
| "", [](tirx::ScopeIdDef def, AccessPath p, IRDocsifier d) -> Doc { |
| auto [parent, cur] = tirx::ScopeBindingToStringPair(def->scope); |
| ExprDoc extents_doc = def->extents.has_value() |
| ? d->AsDoc<ExprDoc>(def->extents.value(), p->Attr("extents")) |
| : LiteralDoc::None(p->Attr("extents")); |
| Doc doc = TIR(d, "ScopeIdDef") |
| ->Call({d->AsDoc<ExprDoc>(def->def_ids, p->Attr("def_ids")), extents_doc, |
| LiteralDoc::Str(parent, p->Attr("parent")), |
| LiteralDoc::Str(cur, p->Attr("cur"))}); |
| return doc; |
| }); |
| TVM_SCRIPT_REPR(tirx::ScopeIdDefNode, ReprPrintTIR); |
| |
| } // namespace printer |
| } // namespace script |
| } // namespace tvm |
