CodegenEnv.h Source File (original) (raw)

13 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENENV_H_

14 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENENV_H_

23 #include

25 namespace mlir {

26 namespace sparse_tensor {

36 public:

41 unsigned numTensors, unsigned numLoops, unsigned maxRank);

50 linalg::GenericOp op() const { return linalgOp; }

55 }

63 std::optional<Operation *>

66 callback);

74 }

77 }

80 }

85 return latticeMerger.getLvlType(t, i);

86 }

103 }

106 }

109 }

110 template

113 }

114

116

117

118

119

120

121

122

124

125

127

128

129

130

131

134

137

140 bool isExpand() const { return expValues != nullptr; }

147

148

149

150

151

157

163

168

169 private:

170

171 linalg::GenericOp linalgOp;

172

173

175

176

177 Merger latticeMerger;

178

179

181

182

183

184

186

193

194

195

196

200

201

202

203

204 Value redValidLexInsert;

205

206

208 };

209

210 }

211 }

212

213 #endif

This class represents an operand of an operation.

This class represents an instance of an SSA value in the MLIR system, representing a computable value...

The code generation environment class aggregates a number of data structures that are needed during t...

void startReduc(ExprId exp, Value val)

void updateValidLexInsert(Value val)

const SparsificationOptions & options() const

Value getInsertionChain() const

std::optional< Operation * > genLoopBoundary(function_ref< std::optional< Operation * >(MutableArrayRef< Value > parameters)> callback)

Generates loop boundary statements (entering/exiting loops).

ArrayRef< LatPointId > set(LatSetId s) const

bool isAdmissibleTensorExp(ExprId e)

Whether the tensor expression is admissible for codegen.

bool atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const

bool isCustomReduc() const

CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts, unsigned numTensors, unsigned numLoops, unsigned maxRank)

Constructs a code generation environment which can be passed around during sparsification for bookkee...

unsigned getCurrentDepth() const

std::pair< TensorId, Level > unpackTensorLevel(TensorLevel tl) const

Value getExpandValues() const

TensorLevel makeTensorLevel(TensorId t, Level l) const

const LatPoint & lat(LatPointId l) const

constexpr TensorId makeTensorId(unsigned t) const

LevelType lt(TensorLoopId b) const

void startExpand(Value values, Value filled, Value added, Value count)

bool hasSparseOutput() const

unsigned getLoopNum() const

void updateInsertionChain(Value chain)

bool generatingSparseIterator() const

Value getExpandCount() const

void startCustomReduc(ExprId exp)

TensorLevel makeTensorLevel(std::pair< TensorId, Level > tlPair) const

constexpr TensorLoopId makeTensorLoopId(unsigned t, unsigned i) const

linalg::GenericOp op() const

Value getLoopVar(LoopId i) const

Returns the induction-variable for the given loop.

Value getExpandFilled() const

LogicalResult initTensorExp()

void startEmit(SparseEmitStrategy emitStrategy)

auto unpackTensorLevelRange(ContainerTy &&c) const

Value getExpandAdded() const

const TensorExp & exp(ExprId e) const

void updateExpandCount(Value count)

void updateReduc(Value val)

Value getValidLexInsert() const

bool isSparseOutput(OpOperand *o) const

void startValidLexInsert(Value val)

constexpr LoopId makeLoopId(unsigned i) const

Value getCustomRedId() const

LevelType lt(TensorId t, LoopId i) const

bool isValidLexInsert() const

TensorId getOutTensorId() const

Gets the TensorId for output tensor.

TensorLevel makeTensorLevel(TensorId t, Level l) const

Compresses a TensorId and Level into a TensorLevel.

unsigned getNumManifestTensors() const

Gets the total number of manifest tensors (excluding the synthetic tensor).

std::pair< TensorId, Level > unpackTensorLevel(TensorLevel tidLvl) const

De-compresses a TensorLevel back to a pair of TensorId and Level.

auto unpackTensorLevelRange(ContainerTy &&c) const

Converts a range of TensorLevel to a range of std::pair<TensorId, Level>

unsigned getNumTensors() const

Gets the total number of tensors that loopEmitter is operating on.

LoopId getCurrentDepth() const

Gets the current depth of the loop-stack.

TensorId getSynTensorId() const

Gets the TensorId for synthetic tensor.

A class to handle all iteration lattice operations.

constexpr unsigned getNumLoops() const

Gets the total number of loops (native loops + filter loops).

constexpr unsigned getNumTensors() const

Gets the total number of tensors (including the output-tensor and synthetic-tensor).

constexpr LoopId makeLoopId(unsigned i) const

Safely converts the argument to a loop identifier.

ArrayRef< LatPointId > set(LatSetId s) const

constexpr TensorId getSynTensorID() const

Gets the synthetic tensor's identifier (used for all invariant tensor expressions).

const TensorExp & exp(ExprId e) const

Convenience getters to immediately access the stored nodes.

const LatPoint & lat(LatPointId p) const

constexpr TensorId getOutTensorID() const

Gets the output tensor's identifier.

constexpr TensorId makeTensorId(unsigned t) const

Safely converts the argument to a tensor identifier.

LevelType getLvlType(TensorId t, LoopId i) const

Gets the level-type of the tth tensor on ith loop.

constexpr TensorLoopId makeTensorLoopId(unsigned t, unsigned i) const

Safely converts the arguments to a pair of (tensor,loop) identifiers.

static constexpr unsigned kInvalidId

A constant serving as the canonically invalid identifier, regardless of the identifier type.

unsigned LatSetId

LatSet identifiers.

unsigned TensorLoopId

A compressed representation of std::pair<TensorId, LoopId>.

uint64_t Level

The type of level identifiers and level-ranks.

unsigned LoopId

Loop identifiers.

unsigned ExprId

TensorExp identifiers.

unsigned LatPointId

LatPoint identifiers.

unsigned TensorId

Tensor identifiers, chosen to be the BlockArgument::getArgNumber of the value passed to Merger::build...

Include the generated interface declarations.

SparseEmitStrategy

Defines a scope for reinterpret map pass.

Options for the Sparsification pass.

SparseEmitStrategy sparseEmitStrategy

This enum defines all the sparse representations supportable by the SparseTensor dialect.

Tensor expression. Represents an MLIR expression in tensor index notation.