102 lines
2.8 KiB
C++
102 lines
2.8 KiB
C++
//===- MapRef.h - A dim2lvl/lvl2dim map encoding ----------------*- C++ -*-===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// A dim2lvl/lvl2dim map encoding class, with utility methods.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_MAPREF_H
|
|
#define MLIR_EXECUTIONENGINE_SPARSETENSOR_MAPREF_H
|
|
|
|
#include <cinttypes>
|
|
|
|
#include <cassert>
|
|
|
|
namespace mlir {
|
|
namespace sparse_tensor {
|
|
|
|
/// A class for capturing the sparse tensor type map with a compact encoding.
|
|
///
|
|
/// Currently, the following situations are supported:
|
|
/// (1) map is a permutation
|
|
/// (2) map has affine ops (restricted set)
|
|
///
|
|
/// The pushforward/backward operations are fast for (1) but incur some obvious
|
|
/// overhead for situation (2).
|
|
///
|
|
class MapRef final {
|
|
public:
|
|
MapRef(uint64_t d, uint64_t l, const uint64_t *d2l, const uint64_t *l2d);
|
|
|
|
//
|
|
// Push forward maps from dimensions to levels.
|
|
//
|
|
|
|
// Map from dimRank in to lvlRank out.
|
|
template <typename T>
|
|
inline void pushforward(const T *in, T *out) const {
|
|
if (isPermutation) {
|
|
for (uint64_t l = 0; l < lvlRank; l++) {
|
|
out[l] = in[dim2lvl[l]];
|
|
}
|
|
} else {
|
|
uint64_t i, c;
|
|
for (uint64_t l = 0; l < lvlRank; l++)
|
|
if (isFloor(l, i, c)) {
|
|
out[l] = in[i] / c;
|
|
} else if (isMod(l, i, c)) {
|
|
out[l] = in[i] % c;
|
|
} else {
|
|
out[l] = in[dim2lvl[l]];
|
|
}
|
|
}
|
|
}
|
|
|
|
//
|
|
// Push backward maps from levels to dimensions.
|
|
//
|
|
|
|
// Map from lvlRank in to dimRank out.
|
|
template <typename T>
|
|
inline void pushbackward(const T *in, T *out) const {
|
|
if (isPermutation) {
|
|
for (uint64_t d = 0; d < dimRank; d++)
|
|
out[d] = in[lvl2dim[d]];
|
|
} else {
|
|
uint64_t i, c, ii;
|
|
for (uint64_t d = 0; d < dimRank; d++)
|
|
if (isMul(d, i, c, ii)) {
|
|
out[d] = in[i] + c * in[ii];
|
|
} else {
|
|
out[d] = in[lvl2dim[d]];
|
|
}
|
|
}
|
|
}
|
|
|
|
uint64_t getDimRank() const { return dimRank; }
|
|
uint64_t getLvlRank() const { return lvlRank; }
|
|
|
|
private:
|
|
bool isPermutationMap() const;
|
|
|
|
bool isFloor(uint64_t l, uint64_t &i, uint64_t &c) const;
|
|
bool isMod(uint64_t l, uint64_t &i, uint64_t &c) const;
|
|
bool isMul(uint64_t d, uint64_t &i, uint64_t &c, uint64_t &ii) const;
|
|
|
|
const uint64_t dimRank;
|
|
const uint64_t lvlRank;
|
|
const uint64_t *const dim2lvl; // non-owning pointer
|
|
const uint64_t *const lvl2dim; // non-owning pointer
|
|
const bool isPermutation;
|
|
};
|
|
|
|
} // namespace sparse_tensor
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_MAPREF_H
|