SparseTensorDescriptor.h Source File (original) (raw)

13 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_SPARSETENSORDESCRIPTOR_H_

14 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_SPARSETENSORDESCRIPTOR_H_

21 namespace mlir {

22 namespace sparse_tensor {

25 public:

27 : specifier(cast<TypedValue>(specifier)) {}

33 operator Value() { return specifier; }

36 StorageSpecifierKind kind, std::optional lvl);

39 StorageSpecifierKind kind, std::optional lvl);

41 private:

43 };

50 template

52 protected:

58 static_assert(std::is_trivially_copyable_v<

60 }

62 public:

64 std::optional lvl) const {

67 }

78 StorageSpecifierKind kind,

79 std::optional lvl) const {

82 }

85 return getSpecifierField(builder, loc, StorageSpecifierKind::LvlSize, lvl);

86 }

90 }

94 }

97 std::optional lvl) const {

99 }

100

102 assert(fidx < fields.size() - 1);

104 }

105

107 return getSpecifierField(builder, loc, StorageSpecifierKind::PosMemSize,

108 lvl);

109 }

110

112 return getSpecifierField(builder, loc, StorageSpecifierKind::CrdMemSize,

113 lvl);

114 }

115

117 return getSpecifierField(builder, loc, StorageSpecifierKind::ValMemSize,

118 std::nullopt);

119 }

120

122 std::optional lvl) const {

124 }

125

127 assert(fidx < fields.size());

129 }

130

132 return fields.drop_back();

133 }

134

137 }

138

143 }

144

148

149 protected:

153 };

154

155

157 public:

160

162 };

163

164

165

166

167

168

169

172 public:

176

177

178

181 }

182

183

184

185

186

187

191 }

192

194 assert(fidx < fields.size() - 1);

196 }

197

199 assert(fidx < fields.size());

201 }

202

204

206 StorageSpecifierKind kind, std::optional lvl,

211 }

212

214 setSpecifierField(builder, loc, StorageSpecifierKind::ValMemSize,

215 std::nullopt, v);

216 }

217

219 setSpecifierField(builder, loc, StorageSpecifierKind::CrdMemSize, lvl, v);

220 }

221

223 setSpecifierField(builder, loc, StorageSpecifierKind::PosMemSize, lvl, v);

224 }

225

227 setSpecifierField(builder, loc, StorageSpecifierKind::LvlSize, lvl, v);

228 }

229 };

230

231

234 return builder.create(loc, TypeRange(tp), values)

235 .getResult(0);

236 }

237

241 }

242

243 inline SparseTensorDescriptor

246 }

247

248 inline MutSparseTensorDescriptor

251 RankedTensorType type) {

252 fields.assign(adaptorValues.begin(), adaptorValues.end());

254 }

255

256 }

257 }

258

259 #endif

union mlir::linalg::@1203::ArityGroupAndKind::Kind kind

This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...

This class helps build Operations.

Operation * create(const OperationState &state)

Creates an operation given the fields represented as an OperationState.

This class provides an abstraction over the various different ranges of value types.

Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...

This class provides an abstraction over the different types of ranges over Values.

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

Using SmallVector for mutable descriptor allows users to reuse it as a tmp buffers to append value fo...

void setMemRefField(SparseTensorFieldKind kind, std::optional< Level > lvl, Value v)

Adds additional setters for mutable descriptor, update the value for required field.

void setMemRefField(FieldIndex fidx, Value v)

MutSparseTensorDescriptor(SparseTensorType stt, SmallVectorImpl< Value > &buffers)

void setSpecifierField(OpBuilder &builder, Location loc, StorageSpecifierKind kind, std::optional< Level > lvl, Value v)

void setSpecifier(Value newSpec)

void setPosMemSize(OpBuilder &builder, Location loc, Level lvl, Value v)

void setValMemSize(OpBuilder &builder, Location loc, Value v)

void setLvlSize(OpBuilder &builder, Location loc, Level lvl, Value v)

void setField(FieldIndex fidx, Value v)

void setCrdMemSize(OpBuilder &builder, Location loc, Level lvl, Value v)

A helper class around an array of values that corresponds to a sparse tensor.

Value getSpecifier() const

Getters: get the value for required field.

RankedTensorType getRankedTensorType() const

ValueArrayRef getFields() const

FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind, std::optional< Level > lvl) const

std::pair< FieldIndex, unsigned > getCrdMemRefIndexAndStride(Level lvl) const

Value getMemRefField(FieldIndex fidx) const

Value getValMemSize(OpBuilder &builder, Location loc) const

Value getSpecifierField(OpBuilder &builder, Location loc, StorageSpecifierKind kind, std::optional< Level > lvl) const

StorageLayout getLayout() const

ValueRange getMemRefFields() const

Value getAOSMemRef() const

Type getMemRefElementType(SparseTensorFieldKind kind, std::optional< Level > lvl) const

unsigned getNumFields() const

Value getCrdMemSize(OpBuilder &builder, Location loc, Level lvl) const

Value getMemRefField(SparseTensorFieldKind kind, std::optional< Level > lvl) const

Value getPosMemSize(OpBuilder &builder, Location loc, Level lvl) const

Value getValMemRef() const

Value getField(FieldIndex fidx) const

Value getLvlSize(OpBuilder &builder, Location loc, Level lvl) const

SparseTensorDescriptorImpl(SparseTensorType stt, ValueArrayRef fields)

Value getPosMemRef(Level lvl) const

Uses ValueRange for immutable descriptors.

Value getCrdMemRefOrView(OpBuilder &builder, Location loc, Level lvl) const

SparseTensorDescriptor(SparseTensorType stt, ValueRange buffers)

void setSpecifierField(OpBuilder &builder, Location loc, Value v, StorageSpecifierKind kind, std::optional< Level > lvl)

SparseTensorSpecifier(Value specifier)

Value getSpecifierField(OpBuilder &builder, Location loc, StorageSpecifierKind kind, std::optional< Level > lvl)

static Value getInitValue(OpBuilder &builder, Location loc, SparseTensorType stt)

A wrapper around RankedTensorType, which has three goals:

Level getLvlRank() const

Returns the level-rank.

Level getAoSCOOStart() const

Returns the starting level of this sparse tensor type for a trailing COO region that spans at least t...

Provides methods to access fields of a sparse tensor with the given encoding.

unsigned getNumFields() const

Gets the total number of fields for the given sparse tensor encoding.

FieldIndex getMemRefFieldIndex(SparseTensorFieldKind kind, std::optional< Level > lvl) const

Gets the field index for required field.

std::pair< FieldIndex, unsigned > getFieldIndexAndStride(SparseTensorFieldKind kind, std::optional< Level > lvl) const

SparseTensorDescriptor getDescriptorFromTensorTuple(ValueRange adaptorValues, RankedTensorType type)

unsigned FieldIndex

The type of field indices.

uint64_t Level

The type of level identifiers and level-ranks.

Value genTuple(OpBuilder &builder, Location loc, Type tp, ValueRange values)

Packs the given values as a "tuple" value.

MemRefType getMemRefType(T &&t)

Convenience method to abbreviate casting getType().

SparseTensorFieldKind

===-------------------------------------------------------------------—===// The sparse tensor storag...

MutSparseTensorDescriptor getMutDescriptorFromTensorTuple(ValueRange adaptorValues, SmallVectorImpl< Value > &fields, RankedTensorType type)

Include the generated interface declarations.

std::conditional_t< std::is_same_v< Ty, mlir::Type >, mlir::Value, detail::TypedValue< Ty > > TypedValue

If Ty is mlir::Type this will select Value instead of having a wrapper around it.