VectorUtils.cpp Source File (original) (raw)

29 #include "llvm/ADT/DenseSet.h"

30 #include "llvm/Support/InterleavedRange.h"

32 #define DEBUG_TYPE "vector-utils"

34 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")

35 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")

37 using namespace mlir;

42 int64_t dim) {

43 if (isa<UnrankedMemRefType, MemRefType>(source.getType()))

44 return b.createOrFold memref::DimOp(loc, source, dim);

45 if (isa<UnrankedTensorType, RankedTensorType>(source.getType()))

46 return b.createOrFold tensor::DimOp(loc, source, dim);

47 llvm_unreachable("Expected MemRefType or TensorType");

48 }

73 for (int64_t permDim : transp) {

74 if (permDim == dim0)

75 return false;

76 if (permDim == dim1)

77 return true;

78 }

80 llvm_unreachable("Ill-formed transpose pattern");

81 }

83 FailureOr<std::pair<int, int>>

85 VectorType srcType = op.getSourceVectorType();

87 for (auto [index, size] : llvm::enumerate(srcType.getShape()))

88 if (size > 1)

89 srcGtOneDims.push_back(index);

91 if (srcGtOneDims.size() != 2)

92 return failure();

98 op.getPermutation()))

99 return failure();

100

101 return std::pair<int, int>(srcGtOneDims[0], srcGtOneDims[1]);

102 }

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

129 if (enclosingLoopToVectorDim.empty())

132 enclosingLoopToVectorDim.begin()->getFirst()->getContext();

135

136 for (auto kvp : enclosingLoopToVectorDim) {

137 assert(kvp.second < perm.size());

139 castaffine::AffineForOp(kvp.first).getInductionVar(), indices);

140 unsigned numIndices = indices.size();

141 unsigned countInvariantIndices = 0;

142 for (unsigned dim = 0; dim < numIndices; ++dim) {

143 if (!invariants.count(indices[dim])) {

145 "permutationMap already has an entry along dim");

147 } else {

148 ++countInvariantIndices;

149 }

150 }

151 assert((countInvariantIndices == numIndices ||

152 countInvariantIndices == numIndices - 1) &&

153 "Vectorization prerequisite violated: at most 1 index may be "

154 "invariant wrt a vectorized loop");

155 (void)countInvariantIndices;

156 }

157 return AffineMap::get(indices.size(), 0, perm, context);

158 }

159

160

161

162

163

164 template

168 while (current) {

169 if ([[maybe_unused]] auto typedParent = dyn_cast(current)) {

170 assert(res.count(current) == 0 && "Already inserted");

171 res.insert(current);

172 }

173 current = current->getParentOp();

174 }

175 return res;

176 }

177

178

180 return getParentsOfTypeaffine::AffineForOp(block);

181 }

182

188 for (auto *forInst : enclosingLoops) {

189 auto it = loopToVectorDim.find(forInst);

190 if (it != loopToVectorDim.end()) {

191 enclosingLoopToVectorDim.insert(*it);

192 }

193 }

195 }

196

201 }

202

203 bool matcher::operatesOnSuperVectorsOf(Operation &op,

204 VectorType subVectorType) {

205

206

207

208

209

210

211

212

213

214 bool mustDivide = false;

215 (void)mustDivide;

216 VectorType superVectorType;

217 if (auto transfer = dyn_cast(op)) {

218 superVectorType = transfer.getVectorType();

219 mustDivide = true;

221 if (!isafunc::ReturnOp(op)) {

222 op.emitError("NYI: assuming only return operations can have 0 "

223 " results at this point");

224 }

225 return false;

227 if (auto v = dyn_cast(op.getResult(0).getType())) {

228 superVectorType = v;

229 } else {

230

231 return false;

232 }

233 } else {

234

235

236 op.emitError("NYI: operation has more than 1 result");

237 return false;

238 }

239

240

241 auto ratio =

242 computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());

243

244

245 assert((ratio || !mustDivide) &&

246 "vector.transfer operation in which super-vector size is not an"

247 " integer multiple of sub-vector size");

248

249

250

251

252

253

254 return ratio.has_value();

255 }

256

258 if (vectorType.isScalable())

259 return false;

260

262 auto vecRank = vectorType.getRank();

263

264 if (!memrefType.areTrailingDimsContiguous(vecRank))

265 return false;

266

267

268 auto memrefShape = memrefType.getShape().take_back(vecRank);

269

270

271

272 auto firstNonMatchingDim =

274 memrefShape.rbegin(), memrefShape.rend());

275 if (firstNonMatchingDim.first == vectorShape.rend())

276 return true;

277

278

279

282

283 return llvm::all_of(leadingDims, [](auto x) { return x == 1; });

284 }

285

286 std::optional

288 if (vType.getRank() <= targetRank)

289 return {};

290

291

292 auto shapeToUnroll = vType.getShape().drop_back(targetRank);

293 auto scalableDimsToUnroll = vType.getScalableDims().drop_back(targetRank);

294 auto it = llvm::find(scalableDimsToUnroll, true);

295 auto firstScalableDim = it - scalableDimsToUnroll.begin();

296 if (firstScalableDim == 0)

297 return {};

298

299 scalableDimsToUnroll = scalableDimsToUnroll.slice(0, firstScalableDim);

300 assert(!llvm::is_contained(scalableDimsToUnroll, true) &&

301 "unexpected leading scalable dimension");

302

303 shapeToUnroll = shapeToUnroll.slice(0, firstScalableDim);

305 }

306

310 auto loc = xfer->getLoc();

311

313 .Casevector::TransferReadOp(

314 [&](auto readOp) { return readOp.getBase(); })

315 .Casevector::TransferWriteOp(

316 [&](auto writeOp) { return writeOp.getOperand(1); });

317

321 return mixedSourceDims;

322 }

323

325 return (type.getRank() > 1) && (type.getNumScalableDims() <= 1);

326 }

327

332 bool useInBoundsInsteadOfMasking) {

333 assert(llvm::none_of(inputVectorSizes,

334 [](int64_t s) { return s == ShapedType::kDynamic; }) &&

335 "invalid input vector sizes");

336 auto sourceShapedType = cast(source.getType());

337 auto sourceShape = sourceShapedType.getShape();

338 assert(sourceShape.size() == inputVectorSizes.size() &&

339 "expected same ranks.");

341 assert(padValue.getType() == sourceShapedType.getElementType() &&

342 "expected same pad element type to match source element type");

343 int64_t readRank = inputVectorSizes.size();

346

347 if (useInBoundsInsteadOfMasking) {

348

349

350 for (unsigned i = 0; i < readRank; i++)

351 inBoundsVal[i] = (sourceShape[i] == inputVectorSizes[i]) &&

352 !ShapedType::isDynamic(sourceShape[i]);

353 }

354 auto transferReadOp = builder.create vector::TransferReadOp(

355 loc,

356 vectorType,

357 source,

359 padValue,

360 inBoundsVal);

361

362 if (llvm::equal(inputVectorSizes, sourceShape) || useInBoundsInsteadOfMasking)

363 return transferReadOp;

366

369 builder.create vector::CreateMaskOp(loc, maskType, mixedSourceDims);

372 }

373

374 LogicalResult

377 LDBG("Iteration space static sizes:" << llvm::interleaved(shape));

378

379 if (inputVectorSizes.size() != shape.size()) {

380 LDBG("Input vector sizes don't match the number of loops");

381 return failure();

382 }

383 if (ShapedType::isDynamicShape(inputVectorSizes)) {

384 LDBG("Input vector sizes can't have dynamic dimensions");

385 return failure();

386 }

387 if (!llvm::all_of(llvm::zip(shape, inputVectorSizes),

388 [](std::tuple<int64_t, int64_t> sizePair) {

389 int64_t staticSize = std::get<0>(sizePair);

390 int64_t inputSize = std::get<1>(sizePair);

391 return ShapedType::isDynamic(staticSize) ||

392 staticSize <= inputSize;

393 })) {

394 LDBG("Input vector sizes must be greater than or equal to iteration space "

395 "static sizes");

396 return failure();

397 }

398 return success();

399 }

static std::optional< VectorShape > vectorShape(Type type)

static bool areDimsTransposedIn2DSlice(int64_t dim0, int64_t dim1, ArrayRef< int64_t > transp)

Given the n-D transpose pattern 'transp', return true if 'dim0' and 'dim1' should be transposed with ...

static SetVector< Operation * > getEnclosingforOps(Block *block)

Returns the enclosing AffineForOp, from closest to farthest.

static AffineMap makePermutationMap(ArrayRef< Value > indices, const DenseMap< Operation *, unsigned > &enclosingLoopToVectorDim)

Constructs a permutation map from memref indices to vector dimension.

static SetVector< Operation * > getParentsOfType(Block *block)

Implementation detail that walks up the parents and records the ones with the specified type.

A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.

static AffineMap get(MLIRContext *context)

Returns a zero result affine map with no dimensions or symbols: () -> ().

Block represents an ordered list of Operations.

Operation * getParentOp()

Returns the closest surrounding operation that contains this block.

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.

This class helps build Operations.

void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)

Create an operation of specific op type at the current insertion point, and immediately try to fold i...

Operation * create(const OperationState &state)

Creates an operation given the fields represented as an OperationState.

Operation is the basic unit of execution within MLIR.

OpResult getResult(unsigned idx)

Get the 'idx'th result of this operation.

Location getLoc()

The source location the operation was defined or derived from.

InFlightDiagnostic emitError(const Twine &message={})

Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...

Block * getBlock()

Returns the operation block that contains this operation.

unsigned getNumResults()

Return the number of results held by this operation.

This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...

A range-style iterator that allows for iterating over the offsets of all potential tiles of size tile...

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

Type getType() const

Return the type of this value.

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

DenseSet< Value, DenseMapInfo< Value > > getInvariantAccesses(Value iv, ArrayRef< Value > indices)

Given an induction variable iv of type AffineForOp and indices of type IndexType, returns the set of ...

bool hasTensorSemantics(Operation *op)

Return "true" if the given op has tensor semantics and should be bufferized.

constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)

SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)

Return the dimensions of the given memref value.

SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)

Return the dimensions of the given tensor value.

bool isContiguousSlice(MemRefType memrefType, VectorType vectorType)

Return true if vectorType is a contiguous slice of memrefType.

LogicalResult isValidMaskedInputVector(ArrayRef< int64_t > shape, ArrayRef< int64_t > inputVectorSizes)

Returns success if inputVectorSizes is a valid masking configuraion for given shape,...

Operation * maskOperation(OpBuilder &builder, Operation *maskableOp, Value mask, Value passthru=Value())

Creates a vector.mask operation around a maskable operation.

FailureOr< std::pair< int, int > > isTranspose2DSlice(vector::TransposeOp op)

Returns two dims that are greater than one if the transposition is applied on a 2D slice.

std::optional< StaticTileOffsetRange > createUnrollIterator(VectorType vType, int64_t targetRank=1)

Returns an iterator for all positions in the leading dimensions of vType up to the targetRank.

Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim)

Helper function that creates a memref::DimOp or tensor::DimOp depending on the type of source.

bool isLinearizableVector(VectorType type)

Returns true if the input Vector type can be linearized.

Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source, ArrayRef< int64_t > inputVectorSizes, Value padValue, bool useInBoundsInsteadOfMasking=false)

Creates a TransferReadOp from source.

SmallVector< OpFoldResult > getMixedSizesXfer(bool hasTensorSemantics, Operation *xfer, RewriterBase &rewriter)

A wrapper for getMixedSizes for vector.transfer_read and vector.transfer_write Ops (for source and de...

Include the generated interface declarations.

AffineExpr getAffineConstantExpr(int64_t constant, MLIRContext *context)

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

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

std::optional< SmallVector< int64_t > > computeShapeRatio(ArrayRef< int64_t > shape, ArrayRef< int64_t > subShape)

Return the multi-dimensional integral ratio of subShape to the trailing dimensions of shape.

AffineExpr getAffineDimExpr(unsigned position, MLIRContext *context)

These free functions allow clients of the API to not use classes in detail.