|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | + |
| 20 | +/*! |
| 21 | + * \file convert_block_to_opaque.cc |
| 22 | + * \brief Convert the blocks to opaque blocks which do not have block vars. |
| 23 | + */ |
| 24 | + |
| 25 | +#include <tvm/tir/stmt_functor.h> |
| 26 | +#include <tvm/tir/transform.h> |
| 27 | + |
| 28 | +#include "../../runtime/thread_storage_scope.h" |
| 29 | +#include "./ir_utils.h" |
| 30 | + |
| 31 | +namespace tvm { |
| 32 | +namespace tir { |
| 33 | + |
| 34 | +std::pair<std::unordered_map<Stmt, std::vector<std::pair<IterVar, Map<String, ObjectRef>>>, |
| 35 | + ObjectPtrHash, ObjectPtrEqual>, |
| 36 | + Map<Var, Var>> |
| 37 | +FindLoopLCA(const Stmt& root) { |
| 38 | + class LCAFinder : public StmtVisitor { |
| 39 | + public: |
| 40 | + void VisitStmt_(const ForNode* op) final { |
| 41 | + stack.push_back(GetRef<Stmt>(op)); |
| 42 | + StmtVisitor::VisitStmt_(op); |
| 43 | + if (op->kind == ForKind::kThreadBinding) { |
| 44 | + UpdateLCA(op); |
| 45 | + } |
| 46 | + stack.pop_back(); |
| 47 | + } |
| 48 | + |
| 49 | + void UpdateLCA(const ForNode* loop) { |
| 50 | + std::string thread_tag = loop->thread_binding.value()->thread_tag; |
| 51 | + { |
| 52 | + Map<String, ObjectRef>* tgt = &annotations[thread_tag]; |
| 53 | + for (const auto& kv : loop->annotations) { |
| 54 | + tgt->Set(kv.first, kv.second); |
| 55 | + } |
| 56 | + } |
| 57 | + IterVar& iter_var = iters[thread_tag]; |
| 58 | + if (!iter_var.defined()) { |
| 59 | + iter_var = IterVar(Range::FromMinExtent(loop->min, loop->extent), // |
| 60 | + loop->loop_var.copy_with_name(thread_tag), // |
| 61 | + loop->thread_binding.value()->iter_type, // |
| 62 | + thread_tag); |
| 63 | + lca[thread_tag] = stack; |
| 64 | + var_subst.Set(loop->loop_var, iter_var->var); |
| 65 | + return; |
| 66 | + } |
| 67 | + var_subst.Set(loop->loop_var, iter_var->var); |
| 68 | + std::vector<Stmt>& path = lca[thread_tag]; |
| 69 | + uint32_t i = 0; |
| 70 | + for (; i < stack.size() && i < path.size(); ++i) { |
| 71 | + if (!stack[i].same_as(path[i])) { |
| 72 | + break; |
| 73 | + } |
| 74 | + } |
| 75 | + path.resize(i); |
| 76 | + } |
| 77 | + |
| 78 | + std::unordered_map<std::string, std::vector<Stmt>> lca; |
| 79 | + std::unordered_map<std::string, IterVar> iters; |
| 80 | + std::unordered_map<std::string, Map<String, ObjectRef>> annotations; |
| 81 | + Map<Var, Var> var_subst; |
| 82 | + std::vector<Stmt> stack; |
| 83 | + }; |
| 84 | + LCAFinder finder; |
| 85 | + finder(root); |
| 86 | + std::unordered_map<Stmt, std::vector<std::pair<IterVar, Map<String, ObjectRef>>>, ObjectPtrHash, |
| 87 | + ObjectPtrEqual> |
| 88 | + result; |
| 89 | + std::vector<std::string> sorted_thread_tags; |
| 90 | + for (const auto& kv : finder.lca) { |
| 91 | + sorted_thread_tags.push_back(kv.first); |
| 92 | + } |
| 93 | + std::sort(sorted_thread_tags.begin(), sorted_thread_tags.end(), |
| 94 | + [](const std::string& lhs, const std::string& rhs) { |
| 95 | + return lhs.size() > rhs.size(); |
| 96 | + runtime::ThreadScope lhs_scope = runtime::ThreadScope::Create(lhs); |
| 97 | + runtime::ThreadScope rhs_scope = runtime::ThreadScope::Create(rhs); |
| 98 | + if (lhs_scope.rank != rhs_scope.rank) { |
| 99 | + return lhs_scope.rank < rhs_scope.rank; |
| 100 | + } |
| 101 | + return lhs_scope.dim_index < rhs_scope.dim_index; |
| 102 | + }); |
| 103 | + for (const auto& thread_tag : sorted_thread_tags) { |
| 104 | + Stmt lca = finder.lca[thread_tag].back(); |
| 105 | + const IterVar& iter = finder.iters[thread_tag]; |
| 106 | + const Map<String, ObjectRef>& annotations = finder.annotations[thread_tag]; |
| 107 | + result[lca].emplace_back(iter, annotations); |
| 108 | + } |
| 109 | + return {result, finder.var_subst}; |
| 110 | +} |
| 111 | + |
| 112 | +/*! |
| 113 | + * \brief Substitute expr via BlockRealize value bindings and convert each block into opaque |
| 114 | + * blocks. |
| 115 | + */ |
| 116 | +class ThreadBindingLifter : public StmtExprMutator { |
| 117 | + public: |
| 118 | + Stmt VisitStmt_(const ForNode* _op) final { |
| 119 | + For op = GetRef<For>(_op); |
| 120 | + bool is_kernel_root = false; |
| 121 | + if (op->kind == ForKind::kThreadBinding) { |
| 122 | + if (iter_lca.empty()) { |
| 123 | + is_kernel_root = true; |
| 124 | + SetKernelRoot(_op); |
| 125 | + } |
| 126 | + } |
| 127 | + For new_op = Downcast<For>(StmtExprMutator::VisitStmt_(_op)); |
| 128 | + Stmt body = std::move(new_op.CopyOnWrite()->body); |
| 129 | + if (auto it = iter_lca.find(op); it != iter_lca.end()) { |
| 130 | + for (const auto& [iter_var, annotation] : it->second) { |
| 131 | + body = For(iter_var->var, iter_var->dom->min, iter_var->dom->extent, |
| 132 | + ForKind::kThreadBinding, std::move(body), |
| 133 | + IterVar(Range(nullptr), Var(iter_var->thread_tag, iter_var->var->dtype), |
| 134 | + kThreadIndex, iter_var->thread_tag), |
| 135 | + annotation); |
| 136 | + } |
| 137 | + } |
| 138 | + if (is_kernel_root) { |
| 139 | + iter_lca.clear(); |
| 140 | + var_subst.clear(); |
| 141 | + } |
| 142 | + if (op->kind == ForKind::kThreadBinding) { |
| 143 | + return body; |
| 144 | + } else { |
| 145 | + new_op.CopyOnWrite()->body = std::move(body); |
| 146 | + return new_op; |
| 147 | + } |
| 148 | + } |
| 149 | + |
| 150 | + void SetKernelRoot(const ForNode* op) { |
| 151 | + auto result = FindLoopLCA(GetRef<Stmt>(op)); |
| 152 | + this->iter_lca = std::move(result.first); |
| 153 | + this->var_subst = std::move(result.second); |
| 154 | + } |
| 155 | + |
| 156 | + PrimExpr VisitExpr_(const VarNode* op) final { |
| 157 | + auto it = var_subst.find(GetRef<Var>(op)); |
| 158 | + if (it != var_subst.end()) { |
| 159 | + return (*it).second; |
| 160 | + } else { |
| 161 | + return GetRef<PrimExpr>(op); |
| 162 | + } |
| 163 | + } |
| 164 | + |
| 165 | + std::unordered_map<Stmt, std::vector<std::pair<IterVar, Map<String, ObjectRef>>>, ObjectPtrHash, |
| 166 | + ObjectPtrEqual> |
| 167 | + iter_lca; |
| 168 | + Map<Var, Var> var_subst; |
| 169 | +}; |
| 170 | + |
| 171 | +PrimFunc LiftThreadBinding(PrimFunc f) { |
| 172 | + // Only apply this pass to TIR that is not from TE schedules |
| 173 | + if (!IsFromLegacyTESchedule(f)) { |
| 174 | + PrimFuncNode* fptr = f.CopyOnWrite(); |
| 175 | + fptr->body = ThreadBindingLifter()(std::move(fptr->body)); |
| 176 | + return f; |
| 177 | + } else { |
| 178 | + return f; |
| 179 | + } |
| 180 | +} |
| 181 | + |
| 182 | +namespace transform { |
| 183 | + |
| 184 | +Pass LiftThreadBinding() { |
| 185 | + auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) { |
| 186 | + return LiftThreadBinding(std::move(f)); |
| 187 | + }; |
| 188 | + return CreatePrimFuncPass(pass_func, 0, "tir.LiftThreadBinding", {}); |
| 189 | +} |
| 190 | + |
| 191 | +TVM_REGISTER_GLOBAL("tir.transform.LiftThreadBinding").set_body_typed(LiftThreadBinding); |
| 192 | +} // namespace transform |
| 193 | + |
| 194 | +} // namespace tir |
| 195 | +} // namespace tvm |
0 commit comments