psrivas2 commented on code in PR #14066: URL: https://github.com/apache/tvm/pull/14066#discussion_r1114437667
########## src/relax/analysis/layout_transformation.cc: ########## @@ -0,0 +1,636 @@ +/* + * 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. + */ + +/*! + * \file relax/analysis/layout_transormation.cc + * \brief Analyze the PrimFunc and suggest layout transformation on it's blocks and buffers based on + * the user provided layout transformations on it's outputs. + */ +#include <tvm/arith/iter_affine_map.h> +#include <tvm/relax/analysis.h> +#include <tvm/tir/stmt_functor.h> + +namespace tvm { +namespace relax { + +using namespace tir; + +using VarSet = std::unordered_set<tir::Var, ObjectPtrHash, ObjectPtrEqual>; +using VarToBlockIndexMap = std::unordered_map<tir::Var, int, ObjectPtrHash, ObjectPtrEqual>; +using VarToIterTypeMap = std::unordered_map<tir::Var, IterVarType, ObjectPtrHash, ObjectPtrEqual>; + +/********** Helper Functions **********/ + +/*! \brief Converts a list to an Array */ +template <typename T> +static inline Array<T> list_to_array(std::list<T> l) { + Array<T> array; + for (auto& v : l) { + array.push_back(v); + } + return array; +} + +/*! \brief Converts an Array to a list */ +template <typename T> +static inline std::list<T> array_to_list(Array<T> array) { + std::list<T> l; + for (auto& v : array) { + l.push_back(v); + } + return l; +} + +/*! \brief Checks if `expr` is dependent on any var in `defs`*/ +static inline bool IsDependentOn(const PrimExpr& expr, const VarSet& vars) { + bool is_dependent = false; + tir::PostOrderVisit(expr, [&](const ObjectRef& node) { + if (const tir::VarNode* op = node.as<tir::VarNode>()) { + if (vars.count(GetRef<tir::Var>(op)) != 0) { + is_dependent = true; + } + } + }); + return is_dependent; +} + +/*! \brief Checks the PrimExpr for any vars not defined in `defs` */ +static inline bool HasUndefinedVars(const PrimExpr& expr, const VarSet& defs) { + bool has_undefined_vars = false; + tir::PostOrderVisit(expr, [&](const ObjectRef& node) { + if (const tir::VarNode* op = node.as<tir::VarNode>()) { + if (defs.count(GetRef<tir::Var>(op)) == 0) { + has_undefined_vars = true; + } + } + }); + return has_undefined_vars; +} + +/* Checks if a transformation is bijective affine over the given ranges*/ +static inline bool IsBijectiveAffine(const IndexMap& m, const Array<Range>& ranges) { + Map<tir::Var, Range> input_iters; + ICHECK_EQ(m->initial_indices.size(), ranges.size()); + for (size_t i = 0; i < ranges.size(); i++) { + input_iters.Set(m->initial_indices[i], ranges[i]); + } + arith::Analyzer analyzer; + auto iter_map_result = DetectIterMap(m->final_indices, input_iters, /* predicate = */ 1, + /*check_level=*/arith::IterMapLevel::Bijective, &analyzer, + /*simplify_trivial_iterators=*/true); + return !iter_map_result->indices.empty(); +} + +/*! \brief Analyzes the indices returned from DetectIterMap analysis. It collects the spatial + * iterators that are used in indices. This is important to get which spatial iter vars are + * accessed in each index of buffer access. + */ +class IndexAnalyzer : public ExprVisitor { + public: + Array<tir::Var> Analyze(const arith::IterSumExpr& expr) { + VisitExpr(expr); + return iterators_; + } + + private: + /*! \brief Override VisitExpr for iter expr type processing */ + void VisitExpr(const PrimExpr& expr) override { + if (const auto* op = expr.as<arith::IterSumExprNode>()) { + for (const auto& arg : op->args) VisitExpr(arg); + VisitExpr(op->base); + return; + } + if (const auto* op = expr.as<arith::IterSplitExprNode>()) { + VisitIterMark(op->source); + VisitExpr(op->lower_factor); + VisitExpr(op->extent); + VisitExpr(op->scale); + return; + } + return ExprVisitor::VisitExpr(expr); + } + + void VisitIterMark(const arith::IterMark& op) { + if (const auto* var = op->source.as<tir::VarNode>()) + iterators_.push_back(GetRef<tir::Var>(var)); + else + VisitExpr(op->source); + VisitExpr(op->extent); + } + + private: + Array<tir::Var> iterators_; +}; + +/*! \brief Analyzes IterMapResult to get the Spatial Layout of buffer access. We define Spatial Review Comment: Done ########## src/relax/analysis/layout_transformation.cc: ########## @@ -0,0 +1,636 @@ +/* + * 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. + */ + +/*! + * \file relax/analysis/layout_transormation.cc + * \brief Analyze the PrimFunc and suggest layout transformation on it's blocks and buffers based on + * the user provided layout transformations on it's outputs. + */ +#include <tvm/arith/iter_affine_map.h> +#include <tvm/relax/analysis.h> +#include <tvm/tir/stmt_functor.h> + +namespace tvm { +namespace relax { + +using namespace tir; + +using VarSet = std::unordered_set<tir::Var, ObjectPtrHash, ObjectPtrEqual>; +using VarToBlockIndexMap = std::unordered_map<tir::Var, int, ObjectPtrHash, ObjectPtrEqual>; +using VarToIterTypeMap = std::unordered_map<tir::Var, IterVarType, ObjectPtrHash, ObjectPtrEqual>; + +/********** Helper Functions **********/ + +/*! \brief Converts a list to an Array */ +template <typename T> +static inline Array<T> list_to_array(std::list<T> l) { + Array<T> array; + for (auto& v : l) { + array.push_back(v); + } + return array; +} + +/*! \brief Converts an Array to a list */ +template <typename T> +static inline std::list<T> array_to_list(Array<T> array) { + std::list<T> l; + for (auto& v : array) { + l.push_back(v); + } + return l; +} + +/*! \brief Checks if `expr` is dependent on any var in `defs`*/ +static inline bool IsDependentOn(const PrimExpr& expr, const VarSet& vars) { + bool is_dependent = false; + tir::PostOrderVisit(expr, [&](const ObjectRef& node) { + if (const tir::VarNode* op = node.as<tir::VarNode>()) { + if (vars.count(GetRef<tir::Var>(op)) != 0) { + is_dependent = true; + } + } + }); + return is_dependent; +} + +/*! \brief Checks the PrimExpr for any vars not defined in `defs` */ +static inline bool HasUndefinedVars(const PrimExpr& expr, const VarSet& defs) { + bool has_undefined_vars = false; + tir::PostOrderVisit(expr, [&](const ObjectRef& node) { + if (const tir::VarNode* op = node.as<tir::VarNode>()) { + if (defs.count(GetRef<tir::Var>(op)) == 0) { + has_undefined_vars = true; + } + } + }); + return has_undefined_vars; +} + +/* Checks if a transformation is bijective affine over the given ranges*/ +static inline bool IsBijectiveAffine(const IndexMap& m, const Array<Range>& ranges) { + Map<tir::Var, Range> input_iters; + ICHECK_EQ(m->initial_indices.size(), ranges.size()); + for (size_t i = 0; i < ranges.size(); i++) { + input_iters.Set(m->initial_indices[i], ranges[i]); + } + arith::Analyzer analyzer; + auto iter_map_result = DetectIterMap(m->final_indices, input_iters, /* predicate = */ 1, + /*check_level=*/arith::IterMapLevel::Bijective, &analyzer, + /*simplify_trivial_iterators=*/true); + return !iter_map_result->indices.empty(); +} + +/*! \brief Analyzes the indices returned from DetectIterMap analysis. It collects the spatial + * iterators that are used in indices. This is important to get which spatial iter vars are + * accessed in each index of buffer access. + */ +class IndexAnalyzer : public ExprVisitor { + public: + Array<tir::Var> Analyze(const arith::IterSumExpr& expr) { + VisitExpr(expr); + return iterators_; + } + + private: + /*! \brief Override VisitExpr for iter expr type processing */ + void VisitExpr(const PrimExpr& expr) override { + if (const auto* op = expr.as<arith::IterSumExprNode>()) { + for (const auto& arg : op->args) VisitExpr(arg); + VisitExpr(op->base); + return; + } + if (const auto* op = expr.as<arith::IterSplitExprNode>()) { + VisitIterMark(op->source); + VisitExpr(op->lower_factor); + VisitExpr(op->extent); + VisitExpr(op->scale); + return; + } + return ExprVisitor::VisitExpr(expr); + } + + void VisitIterMark(const arith::IterMark& op) { + if (const auto* var = op->source.as<tir::VarNode>()) + iterators_.push_back(GetRef<tir::Var>(var)); + else + VisitExpr(op->source); + VisitExpr(op->extent); + } + + private: + Array<tir::Var> iterators_; +}; + +/*! \brief Analyzes IterMapResult to get the Spatial Layout of buffer access. We define Spatial + * Layout of a buffer access as an array of length equal to the dimensions of the buffer. i-th + * element of Spatial Layout contains spatial iter var used from the block iteration domain. For + * indices, where no spatial iter vars are used, the spatial layout element is empty. If any of the + * buffer access indices use multiple spatial iter vars, the spatial layout is undefined. + * Here are a few examples of inferred spatial layout from buffer access. si denotes i-th spatial + * iter var, and ri denotes i-th reduction iter var. + * SpatialLayout(A[s0*constant, s1]) = {s0, s1} + * SpatialLayout(A[s0, constant, r0, s1]) = {s0, null, null, s1} + * SpatialLayout(A[s0 * c + s1]) = undefined + */ +using SpatialLayout = Array<Optional<tir::Var>>; +static inline SpatialLayout GetSpatialLayout(const arith::IterMapResult& iter_map_result) { + ICHECK(!iter_map_result->indices.empty()); + SpatialLayout result; + for (const arith::IterSumExpr& index : iter_map_result->indices) { + IndexAnalyzer index_analyzer; + Array<tir::Var> iter_vars = index_analyzer.Analyze(index); + if (iter_vars.size() >= 2) { + LOG(WARNING) << "[LayoutInference] Unable to get spatial layout of access: " + << arith::NormalizeIterMapToExpr(index); + return {}; + } + if (iter_vars.empty()) { + result.push_back({}); + continue; + } + result.push_back(iter_vars[0]); + } + return result; +} + +/*! \brief Checks if the two spatial layouts are identical. Two empty spatial layouts are treated as Review Comment: Done -- 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]
