// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

/*!
 * \file layout/utils.h
 * \brief Some arith tools for layout & fragment inference
 *
 */

#ifndef TVM_TL_LAYOUT_UTILS_H_
#define TVM_TL_LAYOUT_UTILS_H_

#include <tvm/arith/iter_affine_map.h>

namespace tvm {
namespace tl {

using namespace tir;

/*!
 * \brief Collect the IterSplit that is not used in expr.
 *
 *  If the expr is (x // 2) and x is in Range(4),
 *  than the result should be (x % 2)
 */
Array<arith::IterSplitExpr> DivideUnusedIterators(const Array<PrimExpr>& exprs,
                                                  const Array<IterVar> input_iters,
                                                  arith::Analyzer* analyzer);

/*!
 * \brief Compress the iterator var, remove the unused part of the var not present in the expr
 *
 *  Returns the compressed IterVar as well as the Updated iter sum expression.
 */
std::pair<PrimExpr, IterVar> CompressIterator(const PrimExpr& expr,
                                              const Array<IterVar> input_iters, const Var& var,
                                              arith::Analyzer* analyzer);

/*!
 * \brief Convert the iter splits returned by DivideUnusedIterators into flattened expression
 *
 */
PrimExpr MakeFlattenedExpression(const Array<arith::IterSplitExpr>& splits);

/*!
 * \brief Convert an Array of IterVar to a Map object
 *
 */
Map<Var, Range> ToVMap(const Array<IterVar>& ivs);

/*!
 * \brief Convert a Map object to an Array of IterVar
 *
 */
Array<IterVar> ToIterVars(const Map<Var, Range>& vmap);

}  // namespace tl
}  // namespace tvm

#endif  // TVM_TL_LAYOUT_UTILS_H_
