CodegenUtils.h Source File (original) (raw)

13 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENUTILS_H_

14 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_CODEGENUTILS_H_

26 namespace mlir {

28 class Location;

32 namespace sparse_tensor {

77 template

79 public:

81 : retTypes(retTypes), params(params), genCall(genCall) {}

86 if (!genCall)

87 return genImplementation(retTypes, params, builder, loc);

90 std::string funcName = getMangledFuncName();

91 ModuleOp module = getParentOpOf(builder);

92 MLIRContext *context = module.getContext();

94 auto func = module.lookupSymbolfunc::FuncOp(result.getAttr());

96 if (!func) {

100 func = builder.create func::FuncOp(

101 loc, funcName,

103 func.setPrivate();

104

105 Block *entryBB = func.addEntryBlock();

108

110 genImplementation(retTypes, args, builder, loc);

111 builder.create func::ReturnOp(loc, result);

112 }

113

114 func::CallOp call = builder.create func::CallOp(loc, func, params);

116 }

117

118 private:

119 template

120 OpTp getParentOpOf(OpBuilder &builder) {

122 }

123

124

125 std::string getMangledFuncName() {

126 return static_cast<SubClass *>(this)->getMangledFuncName();

127 }

128

129

130 SmallVector genImplementation(TypeRange retTypes, ValueRange params,

131 OpBuilder &builder, Location loc) {

132 return static_cast<SubClass *>(this)->genImplementation(retTypes, params,

133 builder, loc);

134 }

135

136 private:

137 TypeRange retTypes;

138 ValueRange params;

139 bool genCall;

140 };

141

142

143 Value genCast(OpBuilder &builder, Location loc, Value value, Type dstTy);

144

145

147 Type dstTp);

148

149

150

151

152 Value genIndexLoad(OpBuilder &builder, Location loc, Value mem, ValueRange s);

153

154

155

156

157

158 TypedAttr getOneAttr(Builder &builder, Type tp);

159

160

161

162

164

165

166

167

169 SmallVectorImpl &dstShape,

170 ArrayRef srcShape, ArrayRef staticDstShape,

171 ArrayRef reassociation);

172

173

174 void reshapeCvs(OpBuilder &builder, Location loc,

175 ArrayRef reassociation,

176 ValueRange srcSizes, ValueRange srcCvs,

177 ValueRange dstSizes, SmallVectorImpl &dstCvs);

178

179

180

181

182

183 FlatSymbolRefAttr getFunc(ModuleOp module, StringRef name, TypeRange resultType,

185

186

187

188 func::CallOp createFuncCall(OpBuilder &builder, Location loc, StringRef name,

189 TypeRange resultType, ValueRange operands,

191

192

193

196

197

198

199

200 Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp);

201

202

203

204

205 Value genAlloca(OpBuilder &builder, Location loc, unsigned sz, Type tp,

206 bool staticShape = false);

207

208

209

211

212

213

214

215 Value allocaBuffer(OpBuilder &builder, Location loc, ValueRange values);

216

217

218

219

220

222 RankedTensorType tensorTp, ValueRange sizes);

223

224

225 void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer);

226

227

228 void sizesFromSrc(OpBuilder &builder, SmallVectorImpl &sizes,

229 Location loc, Value src);

230

231

232 Operation *getTop(Operation *op);

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

251 OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order,

252 function_ref<void(ArrayRef, Value)> callback);

253

254

255

256

257

258 SmallVector loadAll(OpBuilder &builder, Location loc, size_t size,

259 Value mem, size_t offsetIdx = 0,

260 Value offsetVal = Value());

261

262

263

264

265

266 void storeAll(OpBuilder &builder, Location loc, Value mem, ValueRange vs,

267 size_t offsetIdx = 0, Value offsetVal = Value());

268

269

270 TypedValue genToMemref(OpBuilder &builder, Location loc,

271 Value tensor);

272

273

274

277

278

279

282

283

284 Value genReader(OpBuilder &builder, Location loc, SparseTensorType stt,

285 Value tensor,

286 SmallVectorImpl &dimSizesValues,

287 Value &dimSizesBuffer);

288

289

290 Value genMapBuffers(OpBuilder &builder, Location loc, SparseTensorType stt,

291 ArrayRef dimSizesValues, Value dimSizesBuffer,

292 SmallVectorImpl &lvlSizesValues,

293 Value &dim2lvlBuffer,

294 Value &lvl2dimBuffer);

295

296

297

298

299

300

301

302

303

304

305

306

307

308

310 if (auto ctp = dyn_cast(tp)) {

311 auto zeroe = builder.getZeroAttr(ctp.getElementType());

312 auto zeroa = builder.getArrayAttr({zeroe, zeroe});

313 return builder.create complex::ConstantOp(loc, tp, zeroa);

314 }

315 return builder.create arith::ConstantOp(loc, tp, builder.getZeroAttr(tp));

316 }

317

318

319

321 if (auto ctp = dyn_cast(tp)) {

322 auto zeroe = builder.getZeroAttr(ctp.getElementType());

323 auto onee = getOneAttr(builder, ctp.getElementType());

324 auto zeroa = builder.getArrayAttr({onee, zeroe});

325 return builder.create complex::ConstantOp(loc, tp, zeroa);

326 }

327 return builder.create arith::ConstantOp(loc, tp, getOneAttr(builder, tp));

328 }

329

330

333 }

334

335

338 }

339

340

343 }

344

345

348 }

349

350

353 }

354

355

358 }

359

360

362 return constantI32(builder, loc, static_cast<uint32_t>(action));

363 }

364

365

367 unsigned width) {

370 }

371

372

373

375 SparseTensorEncodingAttr enc) {

377 }

378

379

380

382 SparseTensorEncodingAttr enc) {

384 }

385

386

388 Type elemTp) {

391 }

392

393

396 return constantI64(builder, loc, static_cast<uint64_t>(lt));

397 }

398

399

401 if (auto complexAttr = dyn_castcomplex::NumberAttr(attr)) {

402 Type tp = cast(complexAttr.getType()).getElementType();

403 return builder.create complex::ConstantOp(

404 loc, complexAttr.getType(),

407 }

408 return builder.create arith::ConstantOp(loc, cast(attr));

409 }

410

411

413 auto rtp = dyn_cast(type);

414 return !rtp || rtp.getRank() == 0;

415 }

416

417 }

418 }

419

420 #endif

Attributes are known-constant values of operations.

Block represents an ordered list of Operations.

Region * getParent() const

Provide a 'getParent' method for ilist_node_with_parent methods.

BlockArgListType getArguments()

This class is a general helper class for creating context-global objects like types,...

TypedAttr getZeroAttr(Type type)

ArrayAttr getArrayAttr(ArrayRef< Attribute > value)

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

MLIRContext is the top-level object for a collection of MLIR operations.

RAII guard to reset the insertion point of the builder when destroyed.

This class helps build Operations.

void setInsertionPointToStart(Block *block)

Sets the insertion point to the start of the specified block.

void setInsertionPoint(Block *block, Block::iterator insertPoint)

Set the insertion point to the specified location.

Operation * create(const OperationState &state)

Creates an operation given the fields represented as an OperationState.

Block * getInsertionBlock() const

Return the block the current insertion point belongs to.

result_range getResults()

ParentT getParentOfType()

Find the first parent operation of the given type, or nullptr if there is no ancestor operation.

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.

type_range getTypes() const

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

Specialization of arith.constant op that returns an integer of index type.

Specialization of arith.constant op that returns an integer value.

A helper class to simplify lowering operations with/without function calls.

FuncCallOrInlineGenerator(TypeRange retTypes, ValueRange params, bool genCall)

SmallVector< Value > genCallOrInline(OpBuilder &builder, Location loc)

@ Type

An inlay hint that for a type annotation.

TypedAttr getOneAttr(Builder &builder, Type tp)

Generates a 1-valued attribute of the given type.

FlatSymbolRefAttr getFunc(ModuleOp module, StringRef name, TypeRange resultType, ValueRange operands, EmitCInterface emitCInterface)

Returns a function reference (first hit also inserts into module).

Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp)

Generates an uninitialized temporary buffer with room for one value of the given type,...

Value constantIndex(OpBuilder &builder, Location loc, int64_t i)

Generates a constant of index type.

void foreachInSparseConstant(OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order, function_ref< void(ArrayRef< Value >, Value)> callback)

Iterate over a sparse constant, generates constantOp for value and coordinates.

Value constantZero(OpBuilder &builder, Location loc, Type tp)

Generates a 0-valued constant of the given type.

Value constantLevelTypeEncoding(OpBuilder &builder, Location loc, LevelType lt)

Generates a constant of the internal dimension level type encoding.

Value allocaBuffer(OpBuilder &builder, Location loc, ValueRange values)

Generates a temporary buffer, initializes it with the given contents, and returns it as type memref<?...

Value constantPosTypeEncoding(OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc)

Generates a constant of the internal type-encoding for position overhead storage.

OverheadType posTypeEncoding(SparseTensorEncodingAttr enc)

Returns the OverheadType for position overhead storage.

Value constantOne(OpBuilder &builder, Location loc, Type tp)

Generates a 1-valued constant of the given type.

OverheadType

Encoding of overhead types (both position overhead and coordinate overhead), for "overloading" @newSp...

Value constantI8(OpBuilder &builder, Location loc, int8_t i)

Generates a constant of i8 type.

Action

The actions performed by @newSparseTensor.

uint64_t Dimension

The type of dimension identifiers and dimension-ranks.

OverheadType crdTypeEncoding(SparseTensorEncodingAttr enc)

Returns the OverheadType for coordinate overhead storage.

TypedValue< BaseMemRefType > genToMemref(OpBuilder &builder, Location loc, Value tensor)

OverheadType overheadTypeEncoding(unsigned width)

Converts an overhead storage bitwidth to its internal type-encoding.

Value constantI1(OpBuilder &builder, Location loc, bool b)

Generates a constant of i1 type.

Value genIndexLoad(OpBuilder &builder, Location loc, Value mem, ValueRange s)

Generates a pointer/index load from the sparse storage scheme.

Value constantAction(OpBuilder &builder, Location loc, Action action)

Generates a constant of the given Action.

Value constantCrdTypeEncoding(OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc)

Generates a constant of the internal type-encoding for coordinate overhead storage.

StringRef overheadTypeFunctionSuffix(OverheadType ot)

Convert OverheadType to its function-name suffix.

PrimaryType

Encoding of the elemental type, for "overloading" @newSparseTensor.

Value genValFromAttr(OpBuilder &builder, Location loc, Attribute attr)

PrimaryType primaryTypeEncoding(Type elemTp)

Converts a primary storage type to its internal type-encoding.

Operation * getTop(Operation *op)

Scans to top of generated loop.

Value createOrFoldSliceStrideOp(OpBuilder &builder, Location loc, Value tensor, Dimension dim)

Generates code to retrieve the slice slice for the sparse tensor slice, return a constant if the offs...

Type getOpaquePointerType(MLIRContext *ctx)

Returns the equivalent of void* for opaque arguments to the execution engine.

Value genMapBuffers(OpBuilder &builder, Location loc, SparseTensorType stt, ArrayRef< Value > dimSizesValues, Value dimSizesBuffer, SmallVectorImpl< Value > &lvlSizesValues, Value &dim2lvlBuffer, Value &lvl2dimBuffer)

Generates code to set up the buffer parameters for a map.

Value genIsNonzero(OpBuilder &builder, Location loc, Value v)

Generates the comparison v != 0 where v is of numeric type.

Value genReader(OpBuilder &builder, Location loc, SparseTensorType stt, Value tensor, SmallVectorImpl< Value > &dimSizesValues, Value &dimSizesBuffer)

Generates code that opens a reader and sets the dimension sizes.

Value genScalarToTensor(OpBuilder &builder, Location loc, Value elem, Type dstTp)

Add conversion from scalar to given type (possibly a 0-rank tensor).

bool isZeroRankedTensorOrScalar(Type type)

void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer)

Generates code to deallocate a dense buffer.

Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp)

Generates an uninitialized temporary buffer of the given size and type, but returns it as type memref...

Value constantOverheadTypeEncoding(OpBuilder &builder, Location loc, unsigned width)

Generates a constant of the internal type-encoding for overhead storage.

SmallVector< Value > loadAll(OpBuilder &builder, Location loc, size_t size, Value mem, size_t offsetIdx=0, Value offsetVal=Value())

Loads size-many values from the memref, which must have rank-1 and size greater-or-equal to size.

void genReshapeDstShape(OpBuilder &builder, Location loc, SmallVectorImpl< Value > &dstShape, ArrayRef< Value > srcShape, ArrayRef< Size > staticDstShape, ArrayRef< ReassociationIndices > reassociation)

Computes the shape of destination tensor of a reshape operator.

Value constantI64(OpBuilder &builder, Location loc, int64_t i)

Generates a constant of i64 type.

void reshapeCvs(OpBuilder &builder, Location loc, ArrayRef< ReassociationIndices > reassociation, ValueRange srcSizes, ValueRange srcCvs, ValueRange dstSizes, SmallVectorImpl< Value > &dstCvs)

Reshape coordinates during a reshaping operation.

Value constantI16(OpBuilder &builder, Location loc, int16_t i)

Generates a constant of i16 type.

func::CallOp createFuncCall(OpBuilder &builder, Location loc, StringRef name, TypeRange resultType, ValueRange operands, EmitCInterface emitCInterface)

Creates a CallOp to the function reference returned by getFunc() in the builder's module.

Value genCast(OpBuilder &builder, Location loc, Value value, Type dstTy)

Add type casting between arith and index types when needed.

StringRef primaryTypeFunctionSuffix(PrimaryType pt)

Convert PrimaryType to its function-name suffix.

Value createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc, Value tensor, Dimension dim)

Generates code to retrieve the slice offset for the sparse tensor slice, return a constant if the off...

Value constantI32(OpBuilder &builder, Location loc, int32_t i)

Generates a constant of i32 type.

Value constantPrimaryTypeEncoding(OpBuilder &builder, Location loc, Type elemTp)

Generates a constant of the internal type-encoding for primary storage.

void sizesFromSrc(OpBuilder &builder, SmallVectorImpl< Value > &sizes, Location loc, Value src)

Populates given sizes array from dense tensor or sparse tensor constant.

Type getOverheadType(Builder &builder, OverheadType ot)

Converts the internal type-encoding for overhead storage to an mlir::Type.

EmitCInterface

Shorthand aliases for the emitCInterface argument to getFunc(), createFuncCall(), and replaceOpWithFu...

Value allocDenseTensor(OpBuilder &builder, Location loc, RankedTensorType tensorTp, ValueRange sizes)

Generates code to allocate a buffer of the given type, and zero initialize it.

void storeAll(OpBuilder &builder, Location loc, Value mem, ValueRange vs, size_t offsetIdx=0, Value offsetVal=Value())

Stores all the values of vs into the memref mem, which must have rank-1 and size greater-or-equal to ...

Include the generated interface declarations.

llvm::function_ref< Fn > function_ref

auto get(MLIRContext *context, Ts &&...params)

Helper method that injects context only if needed, this helps unify some of the attribute constructio...

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