masahi commented on code in PR #14066: URL: https://github.com/apache/tvm/pull/14066#discussion_r1114996824
########## src/relax/analysis/layout_transformation.cc: ########## @@ -0,0 +1,621 @@ +/* + * 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/analysis.h> +#include <tvm/tir/stmt_functor.h> + +#include "../../support/array.h" + +namespace tvm { +namespace relax { + +using namespace tir; + +/********** Helper Functions **********/ + +/*! \brief Checks if a transformation is bijective affine over the given ranges */ +static 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 Analyzer to collect iterators from IterSumExpr. + * \details Analyzes the indices from DetectIterMap analysis to collect the spatial iterators that + * are used in it. This is important to get which spatial iterators 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. + * \details 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 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; + } Review Comment: ``` if (iter_vars.empty()) { result.push_back({}); } else { result.push_back(iter_vars[0]); } ``` -- 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]
