NiccoloN
118 lines
3.6 KiB
C++
118 lines
3.6 KiB
C++
#pragma once
|
|
|
|
#include "mlir/IR/Builders.h"
|
|
|
|
#include "llvm/ADT/ArrayRef.h"
|
|
#include "llvm/ADT/FunctionExtras.h"
|
|
#include "llvm/ADT/Hashing.h"
|
|
#include "llvm/ADT/SmallVector.h"
|
|
#include "llvm/ADT/StringRef.h"
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <optional>
|
|
|
|
namespace onnx_mlir {
|
|
|
|
class ConstantPool;
|
|
|
|
enum class StaticIntSequenceKind {
|
|
Uniform,
|
|
Affine,
|
|
RunLengthEncoded,
|
|
Dense
|
|
};
|
|
|
|
class StaticIntSequence {
|
|
public:
|
|
static StaticIntSequence fromValues(llvm::ArrayRef<int64_t> values);
|
|
static StaticIntSequence uniform(int64_t value, size_t count);
|
|
static StaticIntSequence affine(int64_t base, int64_t step, size_t count);
|
|
|
|
size_t size() const { return count; }
|
|
int64_t valueAt(size_t index) const;
|
|
std::optional<size_t> find(int64_t value, size_t begin, size_t length) const;
|
|
StaticIntSequence slice(size_t begin, size_t count) const;
|
|
StaticIntSequence remap(llvm::ArrayRef<unsigned> indices) const;
|
|
StaticIntSequenceKind getKind() const { return kind; }
|
|
|
|
bool operator==(const StaticIntSequence& other) const;
|
|
llvm::hash_code hash() const;
|
|
|
|
void forEachEqualRun(
|
|
llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
|
|
|
|
private:
|
|
friend class StaticIntSequenceChain;
|
|
friend void setStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef,
|
|
const StaticIntSequence &, size_t);
|
|
friend mlir::FailureOr<StaticIntSequence>
|
|
getStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef, size_t);
|
|
|
|
static StaticIntSequence runLengthEncoded(
|
|
llvm::ArrayRef<int64_t> runs, size_t count);
|
|
StaticIntSequenceKind kind = StaticIntSequenceKind::Dense;
|
|
size_t count = 0;
|
|
int64_t base = 0;
|
|
int64_t step = 0;
|
|
llvm::SmallVector<int64_t> data;
|
|
};
|
|
|
|
struct StaticIntSequenceSlice {
|
|
const StaticIntSequence *sequence = nullptr;
|
|
size_t begin = 0;
|
|
size_t count = 0;
|
|
};
|
|
|
|
class StaticIntSequenceChain {
|
|
public:
|
|
void append(const StaticIntSequence &sequence, size_t begin, size_t count);
|
|
void append(StaticIntSequence sequence);
|
|
size_t size() const { return count; }
|
|
int64_t valueAt(size_t index) const;
|
|
void forEachSegment(llvm::function_ref<void(
|
|
const StaticIntSequence &, size_t, size_t)> callback) const;
|
|
void forEachEqualRun(
|
|
llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
|
|
StaticIntSequence canonicalize() const;
|
|
|
|
private:
|
|
friend class StaticIntSequenceChainCursor;
|
|
llvm::SmallVector<StaticIntSequenceSlice> slices;
|
|
llvm::SmallVector<std::unique_ptr<StaticIntSequence>> owned;
|
|
size_t count = 0;
|
|
};
|
|
|
|
class StaticIntSequenceChainCursor {
|
|
public:
|
|
explicit StaticIntSequenceChainCursor(const StaticIntSequenceChain &chain)
|
|
: chain(chain) {}
|
|
|
|
bool done() const { return slice == chain.slices.size(); }
|
|
int64_t value() const;
|
|
void advance();
|
|
|
|
private:
|
|
const StaticIntSequenceChain &chain;
|
|
size_t slice = 0;
|
|
size_t offset = 0;
|
|
};
|
|
|
|
void setStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
|
|
const StaticIntSequence &sequence,
|
|
size_t logicalCount);
|
|
|
|
mlir::FailureOr<StaticIntSequence>
|
|
getStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
|
|
size_t logicalCount);
|
|
|
|
mlir::Value emitStaticIntLookup(const StaticIntSequence& sequence,
|
|
mlir::Value position,
|
|
mlir::Operation* constantAnchor,
|
|
ConstantPool& constants,
|
|
mlir::OpBuilder& builder,
|
|
mlir::Location loc);
|
|
|
|
} // namespace onnx_mlir
|