junrushao1994 commented on a change in pull request #8467: URL: https://github.com/apache/tvm/pull/8467#discussion_r671922088
########## File path: src/tir/schedule/primitive/loop_transformation.cc ########## @@ -0,0 +1,420 @@ +/* + * 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 tir { + +/*! \brief Append a new predicate to the each child of type BlockRealize (not recursively) */ +class BlockPredicateAppender : public StmtMutator { + public: + /*! + * \brief Constructor + * \param to_append The predicate to be appended to BlockRealizeNode + */ + explicit BlockPredicateAppender(const PrimExpr& to_append, arith::Analyzer* analyzer) + : to_append_(to_append) { + add_predicate_ = !analyzer->CanProve(to_append); + } + + private: + // For each direct child of type BlockRealizeNode, append the predicate + Stmt VisitStmt_(const BlockRealizeNode* realize) final { + // We do not recursively do this + if (add_predicate_) { + ObjectPtr<BlockRealizeNode> n = CopyOnWrite(realize); + n->predicate = n->predicate && to_append_; + return BlockRealize(n); + } else { + return GetRef<BlockRealize>(realize); + } + } + + /*! \brief The predicate to be appended */ + const PrimExpr& to_append_; + /*! \brief Whether to add predicate */ + bool add_predicate_; +}; + +/*! \brief Substitute vars and collect the reuse mapping of opaque blocks */ +class SubstituteVarAndCollectOpaqueBlock : public StmtExprMutator { + public: + explicit SubstituteVarAndCollectOpaqueBlock(std::function<Optional<PrimExpr>(const Var&)> vmap, + Map<Block, Block>* opaque_blocks) + : vmap_(vmap), opaque_blocks_(opaque_blocks) {} + + private: + PrimExpr VisitExpr_(const VarNode* op) final { + Var var = GetRef<Var>(op); + if (Optional<PrimExpr> ret = vmap_(var)) { + return ret.value(); + } else { + return std::move(var); + } + } + + Stmt VisitStmt_(const BlockRealizeNode* op) final { + Stmt res = StmtMutator::VisitStmt_(op); + if (op->block->iter_vars.empty()) { + const BlockRealizeNode* realize = TVM_TYPE_AS(realize, res, BlockRealizeNode); + opaque_blocks_->Set(op->block, realize->block); + } + return res; + } + + /*! \brief The substitute function */ + std::function<Optional<PrimExpr>(const Var&)> vmap_; + /*! \brief The reuse mapping of opaque blocks */ + Map<Block, Block>* opaque_blocks_; +}; + +/*! \brief Simplify the binding of block realize and update the opaque block reuse mapping */ +class IterMapSimplifyBlockBinding : public StmtExprMutator { + public: + explicit IterMapSimplifyBlockBinding(const Map<Var, Range>& loop_map, + Map<Block, Block>* opaque_blocks) + : opaque_blocks_(opaque_blocks), loop_var2extent_(std::move(loop_map)) {} Review comment: 1. Rename `loop_map` to `loop_var2extent` be consistent with the class member 2. Dont need the constant reference, use moving semantics to keep the reference counter unique ```suggestion explicit IterMapSimplifyBlockBinding(Map<Var, Range> loop_var2extent, Map<Block, Block>* opaque_blocks) : opaque_blocks_(opaque_blocks), loop_var2extent_(loop_var2extent) {} ``` -- 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]
