vinx13 commented on a change in pull request #9871: URL: https://github.com/apache/tvm/pull/9871#discussion_r788239076
########## File path: src/tir/schedule/primitive/blockize_tensorize.cc ########## @@ -0,0 +1,638 @@ +/* + * 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 <functional> + +#include "../utils.h" + +namespace tvm { +namespace tir { + +/*! + * \brief ScheduleError that the bindings of the inner block are not divisible by the subspace + * represented by the outer loops. + */ +class SubspaceNotDivisibleError : public ScheduleError { + public: + explicit SubspaceNotDivisibleError(IRModule mod, For scope_loop, Block inner_block) + : mod_(std::move(mod)), + scope_loop_(std::move(scope_loop)), + inner_block_(std::move(inner_block)) {} + + String FastErrorString() const final { + return "ScheduleError: The bindings of the inner block can not be blockized."; + } + + String DetailRenderTemplate() const final { + return "ScheduleError: The bindings of the inner block {0} can not be blockized by the loops " + "starting at {1}."; + } + + IRModule mod() const final { return mod_; } + + Array<ObjectRef> LocationsOfInterest() const final { return {inner_block_, scope_loop_}; } + + private: + IRModule mod_; + For scope_loop_; + Block inner_block_; +}; + +/*! + * \brief Detect if bindings are a trivial case of the subspace division where we can divide the + * block iter bindings into two categories: + * 1. The binding covers no inner loop vars. + * 2. The binding covers only inner loop vars. + * + * The bindings are not required to be quasi-affine. + * + * \param iter_vars The input iterators + * \param bindings The values of iter_vars + * \param outer_loops Iterators outside the subspace. + * \param inner_loops Iterators of the subspace + * \param predicate The predicate constaints on the input iterators. + * \return The result of the subspace division. + */ +Array<Array<arith::IterMark>> TrivialSubspaceDivision(const Array<IterVar>& iter_vars, + const Array<PrimExpr>& bindings, + const Array<Var>& outer_iters, + const Array<Var>& inner_iters, + const PrimExpr& predicate) { + if (!is_one(predicate)) return {}; + std::vector<Array<arith::IterMark>> res; + std::unordered_set<const VarNode*> outer_loop_vars; + std::unordered_set<const VarNode*> inner_loop_vars; + for (const Var& var : outer_iters) { + outer_loop_vars.insert(var.get()); + } + for (const Var& var : inner_iters) { + inner_loop_vars.insert(var.get()); + } + const arith::IterMark unit_iter_mark(arith::IterSumExpr({}, 0), 1); + + for (size_t i = 0; i < bindings.size(); ++i) { + bool outer = UsesVar( + bindings[i], [&outer_loop_vars](const VarNode* var) { return outer_loop_vars.count(var); }); + bool inner = UsesVar( + bindings[i], [&inner_loop_vars](const VarNode* var) { return inner_loop_vars.count(var); }); + arith::IterMark iter_mark; + if (bindings[i]->IsInstance<VarNode>()) { + iter_mark = arith::IterMark( + arith::IterSplitExpr(arith::IterMark(bindings[i], iter_vars[i]->dom->extent)), + iter_vars[i]->dom->extent); + } else { + iter_mark = arith::IterMark(arith::IterSumExpr({}, bindings[i]), iter_vars[i]->dom->extent); + } + if (outer && !inner) { + arith::IterMark outer{nullptr}; + const auto& outer_iter = iter_mark; + const auto& inner_iter = unit_iter_mark; + res.push_back({outer_iter, inner_iter}); + } else if (inner && !outer) { + const auto& outer_iter = unit_iter_mark; + const auto& inner_iter = iter_mark; + res.push_back({outer_iter, inner_iter}); + } else if (!outer && !inner) { + const auto& outer_iter = unit_iter_mark; + const auto& inner_iter = unit_iter_mark; + res.push_back({outer_iter, inner_iter}); + } else { + return {}; + } + } + res.push_back({arith::IterMark(arith::IterSumExpr({}, 0), Bool(true)), + arith::IterMark(arith::IterSumExpr({}, 0), Bool(true))}); Review comment: this is how `SubspaceDivide` is implemented in affine analysis https://github.com/apache/tvm/blob/main/include/tvm/arith/iter_affine_map.h#L344 the last element of the array is the predicates, the extent here has a different meaning -- 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]
