wehub-resource-sync
288 lines
11 KiB
C++
288 lines
11 KiB
C++
/*
|
|
* 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 tvm/tirx/index_map.h
|
|
* \brief Defines a remapping of buffer indices
|
|
*
|
|
* For use with tvm::tirx::Buffer.
|
|
*/
|
|
#ifndef TVM_TIR_INDEX_MAP_H_
|
|
#define TVM_TIR_INDEX_MAP_H_
|
|
|
|
#include <tvm/ffi/container/array.h>
|
|
#include <tvm/ir/expr.h>
|
|
#include <tvm/runtime/tensor.h>
|
|
#include <tvm/tirx/var.h>
|
|
|
|
#include <utility>
|
|
|
|
namespace tvm {
|
|
namespace arith {
|
|
class Analyzer;
|
|
} // namespace arith
|
|
} // namespace tvm
|
|
|
|
namespace tvm {
|
|
namespace tirx {
|
|
|
|
/*!
|
|
* \brief Defines a mapping between two representations of indices
|
|
* into a buffer.
|
|
*
|
|
* This is primarily used for layout transformations of Buffer
|
|
* objects.
|
|
*/
|
|
class IndexMapNode : public ffi::Object {
|
|
public:
|
|
/*! \brief Variables representing the indices prior to remapping.
|
|
*
|
|
* If initial_indices is empty, then final_indices should also be
|
|
* empty, and no mapping is applied.
|
|
*/
|
|
ffi::Array<PrimVar> initial_indices;
|
|
|
|
/*!
|
|
* \brief Expressions defining the indices after remapping.
|
|
*
|
|
* These expressions should only be in terms of the initial_indices,
|
|
* and must be expressible as an IterSumExpr. The mapping from
|
|
* initial_indices to final_indices must be injective.
|
|
*
|
|
* If final_indices is empty, then initial_indices should also be
|
|
* empty, and the map is an identity function.
|
|
*/
|
|
ffi::Array<PrimExpr> final_indices;
|
|
|
|
/*!
|
|
* \brief The inverse index map.
|
|
*
|
|
* When this is defined, IndexMap::Inverse will return the
|
|
* pre-defined inverse index map. Otherwise, the inverse index map
|
|
* will be computed on the fly. It is the user's responsibility to
|
|
* ensure the correctness of the pre-defined inverse index map.
|
|
*
|
|
* \note ffi::ObjectRef is used here instead of IndexMap to avoid circular reference.
|
|
*/
|
|
ffi::Optional<ffi::ObjectRef> inverse_index_map;
|
|
|
|
/*!
|
|
* \brief Default constructor
|
|
*
|
|
* Defines the mapping as an identity function, with initial_indices
|
|
* equal to the final indices.
|
|
*/
|
|
IndexMapNode() {}
|
|
|
|
/*!
|
|
* \brief Map indices to the output space using a fresh analyzer.
|
|
*
|
|
* \param indices The indices in the input space. Should contain
|
|
* one value for each variable in `initial_indices`.
|
|
* \returns The indices in the output space. Contains one value for
|
|
* each expression in `final_indices`.
|
|
*/
|
|
ffi::Array<PrimExpr> MapIndices(const ffi::Array<PrimExpr>& indices) const;
|
|
/*!
|
|
* \brief Map indices to the output space using an existing analyzer.
|
|
*
|
|
* \param indices The indices in the input space. Should contain
|
|
* one value for each variable in `initial_indices`.
|
|
* \param analyzer An analyzer to be used to simplify the resulting expressions.
|
|
* \returns The indices in the output space. Contains one value for
|
|
* each expression in `final_indices`.
|
|
*/
|
|
ffi::Array<PrimExpr> MapIndices(const ffi::Array<PrimExpr>& indices,
|
|
const arith::Analyzer& analyzer) const;
|
|
|
|
/*! \brief Map a memory range to the output space using a fresh analyzer.
|
|
*
|
|
* If contiguous memory locations in the input space are not
|
|
* necessarily contiguous in the output space (e.g. `lambda i:
|
|
* [8*(i%8) + (i//8)]`), then this will return the smallest range
|
|
* such that all valid indices are contained within the given range.
|
|
*
|
|
* \param ranges The ranges in the input space. Should contain one
|
|
* value for each variable in `initial_indices`.
|
|
* \returns The ranges in the output space. Contains one value for
|
|
* each expression in `final_indices`.
|
|
*/
|
|
ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges) const;
|
|
/*! \brief Map a memory range to the output space using an existing analyzer.
|
|
*
|
|
* If contiguous memory locations in the input space are not
|
|
* necessarily contiguous in the output space (e.g. `lambda i:
|
|
* [8*(i%8) + (i//8)]`), then this will return the smallest range
|
|
* such that all valid indices are contained within the given range.
|
|
*
|
|
* \param ranges The ranges in the input space. Should contain one
|
|
* value for each variable in `initial_indices`.
|
|
* \param analyzer An analyzer to be used to simplify the resulting expressions.
|
|
* \returns The ranges in the output space. Contains one value for
|
|
* each expression in `final_indices`.
|
|
*/
|
|
ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges,
|
|
const arith::Analyzer& analyzer) const;
|
|
|
|
/*! \brief Map a buffer shape to the output space using a fresh analyzer.
|
|
*
|
|
* \param shape The buffer shape in the input space. Should contain
|
|
* one value for each variable in `initial_indices`.
|
|
* \returns The buffer shape in the output space. Contains one
|
|
* value for each expression in `final_indices`.
|
|
*/
|
|
ffi::Array<PrimExpr> MapShape(const ffi::Array<PrimExpr>& shape) const;
|
|
/*! \brief Map a buffer shape to the output space using an existing analyzer.
|
|
*
|
|
* \param shape The buffer shape in the input space. Should contain
|
|
* one value for each variable in `initial_indices`.
|
|
* \param analyzer An analyzer to be used to simplify the resulting expressions.
|
|
* \returns The buffer shape in the output space. Contains one
|
|
* value for each expression in `final_indices`.
|
|
*/
|
|
ffi::Array<PrimExpr> MapShape(const ffi::Array<PrimExpr>& shape,
|
|
const arith::Analyzer& analyzer) const;
|
|
|
|
/* \brief Map an Tensor according to this index map
|
|
*
|
|
* \param arr_src The Tensor whose layout is transformed by this index map.
|
|
*
|
|
* \returns The transformed Tensor.
|
|
*/
|
|
runtime::Tensor MapTensor(runtime::Tensor arr_src) const;
|
|
|
|
/*!
|
|
* \brief Convert to string representation in Python.
|
|
* \param f_name_map Optional function to specify the stringified name of the variables.
|
|
* \return The stringified lambda expression in Python.
|
|
*/
|
|
ffi::String ToPythonString(
|
|
const std::function<ffi::Optional<ffi::String>(const Var& var)>& f_name_map = nullptr) const;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<IndexMapNode>()
|
|
.def_ro("initial_indices", &IndexMapNode::initial_indices,
|
|
refl::AttachFieldFlag::SEqHashDefRecursive())
|
|
.def_ro("final_indices", &IndexMapNode::final_indices)
|
|
.def_ro("inverse_index_map", &IndexMapNode::inverse_index_map,
|
|
refl::AttachFieldFlag::SEqHashIgnore());
|
|
}
|
|
|
|
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.IndexMap", IndexMapNode, ffi::Object);
|
|
};
|
|
|
|
class IndexMap : public ffi::ObjectRef {
|
|
public:
|
|
/*!
|
|
* \brief The constructor
|
|
* \param initial_indices Variables representing the indices prior to remapping
|
|
* \param final_indices Expressions defining the indices after remapping.
|
|
* \param inverse_index_map The optional pre-defined inverse index map
|
|
*/
|
|
IndexMap(ffi::Array<PrimVar> initial_indices, ffi::Array<PrimExpr> final_indices,
|
|
ffi::Optional<IndexMap> inverse_index_map = std::nullopt);
|
|
|
|
/*!
|
|
* \brief Create an index map from a packed function
|
|
* \param ndim The number of dimensions
|
|
* \param func The function to be applied
|
|
* \param inverse_index_map The optional pre-defined inverse index map
|
|
* \return The created index map
|
|
*/
|
|
static IndexMap FromFunc(int ndim, ffi::TypedFunction<ffi::Array<PrimExpr>(ffi::Array<Var>)> func,
|
|
ffi::Optional<IndexMap> inverse_index_map = std::nullopt);
|
|
|
|
/*! \brief Generate the inverse mapping using a fresh analyzer.
|
|
*
|
|
* The range of the input indices is required in order to ensure
|
|
* that the transformation is bijective over the input domain.
|
|
*
|
|
* If the user has supplied an `inverse_index_map`, that map is
|
|
* assumed to be correct and bijective, and is returned.
|
|
* \param initial_ranges The ranges of the input indices.
|
|
*/
|
|
IndexMap Inverse(ffi::Array<Range> initial_ranges) const;
|
|
/*! \brief Generate the inverse mapping using an existing analyzer.
|
|
*
|
|
* The range of the input indices is required in order to ensure
|
|
* that the transformation is bijective over the input domain.
|
|
*
|
|
* If the user has supplied an `inverse_index_map`, that map is
|
|
* assumed to be correct and bijective, and is returned.
|
|
* \param initial_ranges The ranges of the input indices.
|
|
* \param analyzer An analyzer to be used while deriving and validating
|
|
* the inverse.
|
|
*/
|
|
IndexMap Inverse(ffi::Array<Range> initial_ranges, const arith::Analyzer& analyzer) const;
|
|
|
|
/*! \brief Rename the variables in the index map and ensure the names are unique.
|
|
*
|
|
* Construct a new index map with the same transformation, but with name_hint of variables to be
|
|
* guaranteed unique. The optional f_name_map can be provided to rename the variables.
|
|
*
|
|
* \param f_name_map The optional name map to rename the variables.
|
|
* \return The renamed index map.
|
|
*/
|
|
IndexMap RenameVariables(
|
|
const std::function<ffi::Optional<ffi::String>(const Var& var)>& f_name_map = nullptr) const;
|
|
|
|
/*! \brief Generate the inverse mapping using a fresh analyzer.
|
|
*
|
|
* Determine the inverse, where the output range may contain
|
|
* addresses that do not correspond to an address in the input
|
|
* range.
|
|
*
|
|
* \param initial_ranges The ranges of the input indices.
|
|
* \return The inverted index map, along with the predicate for
|
|
* which the inverse maps to a valid range.
|
|
*/
|
|
std::pair<IndexMap, PrimExpr> NonSurjectiveInverse(ffi::Array<Range> initial_ranges) const;
|
|
/*! \brief Generate the inverse mapping using an existing analyzer.
|
|
*
|
|
* Determine the inverse, where the output range may contain
|
|
* addresses that do not correspond to an address in the input
|
|
* range.
|
|
*
|
|
* \param initial_ranges The ranges of the input indices.
|
|
* \param analyzer An analyzer to be used while deriving the inverse and
|
|
* padding predicate.
|
|
* \return The inverted index map, along with the predicate for
|
|
* which the inverse maps to a valid range.
|
|
*/
|
|
std::pair<IndexMap, PrimExpr> NonSurjectiveInverse(ffi::Array<Range> initial_ranges,
|
|
const arith::Analyzer& analyzer) const;
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IndexMap, ffi::ObjectRef, IndexMapNode);
|
|
};
|
|
|
|
/*! \brief Substitute variables in an index map.
|
|
*
|
|
* \param index_map The index_map
|
|
* \param f_subst The substitution function
|
|
*/
|
|
IndexMap Substitute(const IndexMap& index_map,
|
|
std::function<ffi::Optional<PrimExpr>(const Var& var)> f_subst);
|
|
|
|
} // namespace tirx
|
|
} // namespace tvm
|
|
|
|
#endif // TVM_TIR_INDEX_MAP_H_
|