1	//===- VectorOps.cpp - MLIR Vector Dialect Operations ---------------------===//
2	//
3	// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4	// See https://llvm.org/LICENSE.txt for license information.
5	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6	//
7	//===----------------------------------------------------------------------===//
8	//
9	// This file implements convenience types for working with super-vectorization
10	// operations, in particular super-vector loads and stores.
11	//
12	//===----------------------------------------------------------------------===//
13
14	#include "mlir/Dialect/Vector/IR/VectorOps.h"
15
16	#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
17	#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
18	#include "mlir/Dialect/Arith/IR/Arith.h"
19	#include "mlir/Dialect/Arith/Utils/Utils.h"
20	#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
21	#include "mlir/Dialect/MemRef/IR/MemRef.h"
22	#include "mlir/Dialect/Tensor/IR/Tensor.h"
23	#include "mlir/Dialect/UB/IR/UBOps.h"
24	#include "mlir/Dialect/Utils/IndexingUtils.h"
25	#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
26	#include "mlir/IR/AffineExpr.h"
27	#include "mlir/IR/AffineMap.h"
28	#include "mlir/IR/Builders.h"
29	#include "mlir/IR/BuiltinAttributes.h"
30	#include "mlir/IR/BuiltinTypes.h"
31	#include "mlir/IR/DialectImplementation.h"
32	#include "mlir/IR/IRMapping.h"
33	#include "mlir/IR/OpImplementation.h"
34	#include "mlir/IR/PatternMatch.h"
35	#include "mlir/IR/TypeUtilities.h"
36	#include "mlir/IR/ValueRange.h"
37	#include "mlir/Interfaces/SubsetOpInterface.h"
38	#include "mlir/Interfaces/ValueBoundsOpInterface.h"
39	#include "mlir/Support/LLVM.h"
40	#include "mlir/Transforms/InliningUtils.h"
41	#include "llvm/ADT/ArrayRef.h"
42	#include "llvm/ADT/STLExtras.h"
43	#include "llvm/ADT/SmallVector.h"
44	#include "llvm/ADT/StringSet.h"
45	#include "llvm/ADT/TypeSwitch.h"
46	#include "llvm/Support/Casting.h"
47
48	#include <cassert>
49	#include <cstdint>
50	#include <numeric>
51
52	#include "mlir/Dialect/Vector/IR/VectorDialect.cpp.inc"
53	// Pull in all enum type and utility function definitions.
54	#include "mlir/Dialect/Vector/IR/VectorEnums.cpp.inc"
55
56	using namespace mlir;
57	using namespace mlir::vector;
58
59	/// Helper enum to classify mask value.
60	enum class MaskFormat {
61	AllTrue = 0,
62	AllFalse = 1,
63	Unknown = 2,
64	};
65
66	/// Helper method to classify a mask value. Currently, the method
67	/// looks "under the hood" of a constant value with dense attributes
68	/// and a constant mask operation (since the client may be called at
69	/// various stages during progressive lowering).
70	static MaskFormat getMaskFormat(Value mask) {
71	if (auto c = mask.getDefiningOp<arith::ConstantOp>()) {
72	// Inspect constant dense values. We count up for bits that
73	// are set, count down for bits that are cleared, and bail
74	// when a mix is detected.
75	if (auto denseElts = llvm::dyn_cast<DenseIntElementsAttr>(c.getValue())) {
76	int64_t val = 0;
77	for (bool b : denseElts.getValues<bool>())
78	if (b && val >= 0)
79	val++;
80	else if (!b && val <= 0)
81	val--;
82	else
83	return MaskFormat::Unknown;
84	if (val > 0)
85	return MaskFormat::AllTrue;
86	if (val < 0)
87	return MaskFormat::AllFalse;
88	}
89	} else if (auto m = mask.getDefiningOp<ConstantMaskOp>()) {
90	// Inspect constant mask index. If the index exceeds the
91	// dimension size, all bits are set. If the index is zero
92	// or less, no bits are set.
93	ArrayRef<int64_t> masks = m.getMaskDimSizes();
94	auto shape = m.getType().getShape();
95	bool allTrue = true;
96	bool allFalse = true;
97	for (auto [maskIdx, dimSize] : llvm::zip_equal(masks, shape)) {
98	if (maskIdx < dimSize)
99	allTrue = false;
100	if (maskIdx > 0)
101	allFalse = false;
102	}
103	if (allTrue)
104	return MaskFormat::AllTrue;
105	if (allFalse)
106	return MaskFormat::AllFalse;
107	} else if (auto m = mask.getDefiningOp<CreateMaskOp>()) {
108	// Finds all-false create_masks. An all-true create_mask requires all
109	// dims to be constants, so that'll be folded to a constant_mask, then
110	// detected in the constant_mask case.
111	auto maskOperands = m.getOperands();
112	for (Value operand : maskOperands) {
113	if (auto constantOp = operand.getDefiningOp<arith::ConstantOp>()) {
114	int64_t dimSize =
115	llvm::cast<IntegerAttr>(constantOp.getValue()).getInt();
116	if (dimSize <= 0)
117	return MaskFormat::AllFalse;
118	}
119	}
120	return MaskFormat::Unknown;
121	}
122	return MaskFormat::Unknown;
123	}
124
125	/// Default callback to build a region with a 'vector.yield' terminator with no
126	/// arguments.
127	void mlir::vector::buildTerminatedBody(OpBuilder &builder, Location loc) {
128	builder.create<vector::YieldOp>(loc);
129	}
130
131	// Helper for verifying combining kinds in contractions and reductions.
132	static bool isSupportedCombiningKind(CombiningKind combiningKind,
133	Type elementType) {
134	switch (combiningKind) {
135	case CombiningKind::ADD:
136	case CombiningKind::MUL:
137	return elementType.isIntOrIndexOrFloat();
138	case CombiningKind::MINUI:
139	case CombiningKind::MINSI:
140	case CombiningKind::MAXUI:
141	case CombiningKind::MAXSI:
142	case CombiningKind::AND:
143	case CombiningKind::OR:
144	case CombiningKind::XOR:
145	return elementType.isIntOrIndex();
146	case CombiningKind::MINNUMF:
147	case CombiningKind::MAXNUMF:
148	case CombiningKind::MINIMUMF:
149	case CombiningKind::MAXIMUMF:
150	return llvm::isa<FloatType>(Val: elementType);
151	}
152	return false;
153	}
154
155	/// Returns the effective rank of the vector to read/write for Xfer Ops
156	///
157	/// When the element type of the shaped type is _a scalar_, this will simply
158	/// return the rank of the vector ( the result for xfer_read or the value to
159	/// store for xfer_write).
160	///
161	/// When the element type of the base shaped type is _a vector_, returns the
162	/// difference between the original vector type and the element type of the
163	/// shaped type.
164	///
165	/// EXAMPLE 1 (element type is _a scalar_):
166	/// - shapedType = tensor<10x20xf32>, vectorType = vector<2x4xf32>
167	/// - shapedType.getElementType() = f32 (rank 0)
168	/// - vectorType.getRank() = 2
169	/// - Result = 2 - 0 = 2
170	///
171	/// EXAMPLE 2 (element type is _a vector_):
172	/// - shapedType = tensor<10xvector<20xf32>>, vectorType = vector<20xf32>
173	/// - shapedType.getElementType() = vector<20xf32> (rank 1)
174	/// - vectorType.getRank() = 1
175	/// - Result = 1 - 1 = 0
176	///
177	/// This is used to determine the number of minor dimensions for identity maps
178	/// in vector transfer Ops.
179	static unsigned getEffectiveVectorRankForXferOp(ShapedType shapedType,
180	VectorType vectorType) {
181	unsigned elementVectorRank = 0;
182	VectorType elementVectorType =
183	llvm::dyn_cast<VectorType>(shapedType.getElementType());
184	if (elementVectorType)
185	elementVectorRank += elementVectorType.getRank();
186	return vectorType.getRank() - elementVectorRank;
187	}
188
189	AffineMap mlir::vector::getTransferMinorIdentityMap(ShapedType shapedType,
190	VectorType vectorType) {
191	// 0-d transfers are to/from tensor<t>/memref<t> and vector<1xt>.
192	// TODO: replace once we have 0-d vectors.
193	if (shapedType.getRank() == 0 &&
194	vectorType.getShape() == ArrayRef<int64_t>{1})
195	return AffineMap::get(
196	/*numDims=*/0, /*numSymbols=*/0,
197	getAffineConstantExpr(0, shapedType.getContext()));
198	return AffineMap::getMinorIdentityMap(
199	shapedType.getRank(),
200	getEffectiveVectorRankForXferOp(shapedType, vectorType),
201	shapedType.getContext());
202	}
203
204	/// Check if `write` is of a constant splat and the masked `read` is padded with
205	/// the same splat value -- meaning it could be the same value as the initial
206	/// constant splat.
207	static bool isSplatWriteConsistentWithMaskedRead(vector::TransferWriteOp write,
208	vector::TransferReadOp read) {
209	auto readMask = read.getMask();
210	auto writeMask = write.getMask();
211	// Check if the masks are consistent. The splat value could be the same if the
212	// read is masked (and padded with the splat value), and the write is unmasked
213	// or has the same mask. Note this does not allow the case where the write is
214	// masked and the read is unmasked, as then the read could be of more elements
215	// than the write (which may not be the same value).
216	bool couldBeSameSplat = readMask && (!writeMask || writeMask == readMask);
217	if (!couldBeSameSplat)
218	return false;
219	// Check for constant splat (as the source of the write).
220	DenseElementsAttr splatAttr;
221	if (!matchPattern(write.getVector(),
222	m_Constant<DenseElementsAttr>(bind_value: &splatAttr)) ||
223	!splatAttr.isSplat()) {
224	return false;
225	}
226	// The padding of the read and the constant splat value must be the same.
227	Attribute padAttr;
228	if (!matchPattern(read.getPadding(), m_Constant(bind_value: &padAttr)))
229	return false;
230	return padAttr == splatAttr.getSplatValue<Attribute>();
231	}
232
233	bool mlir::vector::checkSameValueRAW(vector::TransferWriteOp defWrite,
234	vector::TransferReadOp read) {
235	return !defWrite.hasOutOfBoundsDim() &&
236	defWrite.getIndices() == read.getIndices() &&
237	defWrite.getVectorType() == read.getVectorType() &&
238	defWrite.getPermutationMap() == read.getPermutationMap() &&
239	((!defWrite.getMask() && !read.getMask()) ||
240	isSplatWriteConsistentWithMaskedRead(defWrite, read));
241	}
242
243	bool mlir::vector::checkSameValueWAW(vector::TransferWriteOp write,
244	vector::TransferWriteOp priorWrite) {
245	return priorWrite.getIndices() == write.getIndices() &&
246	priorWrite.getMask() == write.getMask() &&
247	priorWrite.getVectorType() == write.getVectorType() &&
248	priorWrite.getPermutationMap() == write.getPermutationMap();
249	}
250
251	bool mlir::vector::isDisjointTransferIndices(
252	VectorTransferOpInterface transferA, VectorTransferOpInterface transferB,
253	bool testDynamicValueUsingBounds) {
254	// For simplicity only look at transfer of same type.
255	if (transferA.getVectorType() != transferB.getVectorType())
256	return false;
257	unsigned rankOffset = transferA.getLeadingShapedRank();
258	for (unsigned i = 0, e = transferA.getIndices().size(); i < e; i++) {
259	Value indexA = transferA.getIndices()[i];
260	Value indexB = transferB.getIndices()[i];
261	std::optional<int64_t> cstIndexA = getConstantIntValue(ofr: indexA);
262	std::optional<int64_t> cstIndexB = getConstantIntValue(ofr: indexB);
263
264	if (i < rankOffset) {
265	// For leading dimensions, if we can prove that index are different we
266	// know we are accessing disjoint slices.
267	if (cstIndexA.has_value() && cstIndexB.has_value()) {
268	if (*cstIndexA != *cstIndexB)
269	return true;
270	continue;
271	}
272	if (testDynamicValueUsingBounds) {
273	// First try to see if we can fully compose and simplify the affine
274	// expression as a fast track.
275	FailureOr<uint64_t> delta =
276	affine::fullyComposeAndComputeConstantDelta(value1: indexA, value2: indexB);
277	if (succeeded(Result: delta) && *delta != 0)
278	return true;
279
280	FailureOr<bool> testEqual =
281	ValueBoundsConstraintSet::areEqual(var1: indexA, var2: indexB);
282	if (succeeded(Result: testEqual) && !testEqual.value())
283	return true;
284	}
285	} else {
286	// For this dimension, we slice a part of the memref we need to make sure
287	// the intervals accessed don't overlap.
288	int64_t vectorDim = transferA.getVectorType().getDimSize(i - rankOffset);
289	if (cstIndexA.has_value() && cstIndexB.has_value()) {
290	int64_t distance = std::abs(i: *cstIndexA - *cstIndexB);
291	if (distance >= vectorDim)
292	return true;
293	continue;
294	}
295	if (testDynamicValueUsingBounds) {
296	// First try to see if we can fully compose and simplify the affine
297	// expression as a fast track.
298	FailureOr<int64_t> delta =
299	affine::fullyComposeAndComputeConstantDelta(value1: indexA, value2: indexB);
300	if (succeeded(Result: delta) && std::abs(i: *delta) >= vectorDim)
301	return true;
302
303	FailureOr<int64_t> computeDelta =
304	ValueBoundsConstraintSet::computeConstantDelta(value1: indexA, value2: indexB);
305	if (succeeded(Result: computeDelta)) {
306	if (std::abs(i: computeDelta.value()) >= vectorDim)
307	return true;
308	}
309	}
310	}
311	}
312	return false;
313	}
314
315	bool mlir::vector::isDisjointTransferSet(VectorTransferOpInterface transferA,
316	VectorTransferOpInterface transferB,
317	bool testDynamicValueUsingBounds) {
318	if (transferA.getBase() != transferB.getBase())
319	return false;
320	return isDisjointTransferIndices(transferA, transferB,
321	testDynamicValueUsingBounds);
322	}
323
324	// Helper to iterate over n-D vector slice elements. Calculate the next
325	// `position` in the n-D vector of size `shape`, applying an offset `offsets`.
326	// Modifies the `position` in place. Returns a failure when `position` becomes
327	// the end position.
328	static LogicalResult incSlicePosition(MutableArrayRef<int64_t> position,
329	ArrayRef<int64_t> shape,
330	ArrayRef<int64_t> offsets) {
331	for (auto [posInDim, dimSize, offsetInDim] :
332	llvm::reverse(C: llvm::zip_equal(t&: position, u&: shape, args&: offsets))) {
333	++posInDim;
334	if (posInDim < dimSize + offsetInDim)
335	return success();
336
337	// Carry the overflow to the next loop iteration.
338	posInDim = offsetInDim;
339	}
340
341	return failure();
342	}
343
344	/// Returns the integer numbers in `values`. `values` are expected to be
345	/// constant operations.
346	SmallVector<int64_t> vector::getAsIntegers(ArrayRef<Value> values) {
347	SmallVector<int64_t> ints;
348	llvm::transform(Range&: values, d_first: std::back_inserter(x&: ints), F: [](Value value) {
349	auto constOp = value.getDefiningOp<arith::ConstantIndexOp>();
350	assert(constOp && "Unexpected non-constant index");
351	return constOp.value();
352	});
353	return ints;
354	}
355
356	/// Returns the integer numbers in `foldResults`. `foldResults` are expected to
357	/// be constant operations.
358	SmallVector<int64_t> vector::getAsIntegers(ArrayRef<OpFoldResult> foldResults) {
359	SmallVector<int64_t> ints;
360	llvm::transform(
361	Range&: foldResults, d_first: std::back_inserter(x&: ints), F: [](OpFoldResult foldResult) {
362	assert(isa<Attribute>(foldResult) && "Unexpected non-constant index");
363	return cast<IntegerAttr>(cast<Attribute>(foldResult)).getInt();
364	});
365	return ints;
366	}
367
368	/// Convert `foldResults` into Values. Integer attributes are converted to
369	/// constant op.
370	SmallVector<Value> vector::getAsValues(OpBuilder &builder, Location loc,
371	ArrayRef<OpFoldResult> foldResults) {
372	SmallVector<Value> values;
373	llvm::transform(Range&: foldResults, d_first: std::back_inserter(x&: values),
374	F: [&](OpFoldResult foldResult) {
375	if (auto attr = dyn_cast<Attribute>(foldResult))
376	return builder
377	.create<arith::ConstantIndexOp>(
378	loc, cast<IntegerAttr>(attr).getInt())
379	.getResult();
380
381	return cast<Value>(foldResult);
382	});
383	return values;
384	}
385
386	std::optional<int64_t> vector::getConstantVscaleMultiplier(Value value) {
387	if (value.getDefiningOp<vector::VectorScaleOp>())
388	return 1;
389	auto mul = value.getDefiningOp<arith::MulIOp>();
390	if (!mul)
391	return {};
392	auto lhs = mul.getLhs();
393	auto rhs = mul.getRhs();
394	if (lhs.getDefiningOp<vector::VectorScaleOp>())
395	return getConstantIntValue(rhs);
396	if (rhs.getDefiningOp<vector::VectorScaleOp>())
397	return getConstantIntValue(lhs);
398	return {};
399	}
400
401	//===----------------------------------------------------------------------===//
402	// CombiningKindAttr
403	//===----------------------------------------------------------------------===//
404
405	namespace mlir {
406	namespace vector {
407	namespace detail {
408	struct BitmaskEnumStorage : public AttributeStorage {
409	using KeyTy = uint64_t;
410
411	BitmaskEnumStorage(KeyTy val) : value(val) {}
412
413	bool operator==(const KeyTy &key) const { return value == key; }
414
415	static BitmaskEnumStorage *construct(AttributeStorageAllocator &allocator,
416	const KeyTy &key) {
417	return new (allocator.allocate<BitmaskEnumStorage>())
418	BitmaskEnumStorage(key);
419	}
420
421	KeyTy value = 0;
422	};
423	} // namespace detail
424	} // namespace vector
425	} // namespace mlir
426
427	//===----------------------------------------------------------------------===//
428	// VectorDialect
429	//===----------------------------------------------------------------------===//
430
431	namespace {
432	/// This class defines the interface for handling inlining with vector dialect
433	/// operations.
434	struct VectorInlinerInterface : public DialectInlinerInterface {
435	using DialectInlinerInterface::DialectInlinerInterface;
436
437	/// All vector dialect ops can be inlined.
438	bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
439	return true;
440	}
441	};
442	} // namespace
443
444	void VectorDialect::initialize() {
445	addAttributes<
446	#define GET_ATTRDEF_LIST
447	#include "mlir/Dialect/Vector/IR/VectorAttributes.cpp.inc"
448	>();
449
450	addOperations<
451	#define GET_OP_LIST
452	#include "mlir/Dialect/Vector/IR/VectorOps.cpp.inc"
453	>();
454
455	addInterfaces<VectorInlinerInterface>();
456
457	declarePromisedInterfaces<bufferization::BufferizableOpInterface,
458	TransferReadOp, TransferWriteOp, GatherOp, MaskOp,
459	YieldOp>();
460	declarePromisedInterfaces<SubsetOpInterface, TransferReadOp,
461	TransferWriteOp>();
462	declarePromisedInterface<SubsetExtractionOpInterface, TransferReadOp>();
463	declarePromisedInterface<SubsetInsertionOpInterface, TransferWriteOp>();
464	declarePromisedInterface<ConvertToLLVMPatternInterface, VectorDialect>();
465	}
466
467	/// Materialize a single constant operation from a given attribute value with
468	/// the desired resultant type.
469	Operation *VectorDialect::materializeConstant(OpBuilder &builder,
470	Attribute value, Type type,
471	Location loc) {
472	if (isa<ub::PoisonAttrInterface>(value))
473	return value.getDialect().materializeConstant(builder, value, type, loc);
474
475	return arith::ConstantOp::materialize(builder, value, type, loc);
476	}
477
478	IntegerType vector::getVectorSubscriptType(Builder &builder) {
479	return builder.getIntegerType(64);
480	}
481
482	ArrayAttr vector::getVectorSubscriptAttr(Builder &builder,
483	ArrayRef<int64_t> values) {
484	return builder.getI64ArrayAttr(values);
485	}
486
487	//===----------------------------------------------------------------------===//
488	// MultiDimReductionOp
489	//===----------------------------------------------------------------------===//
490
491	void vector::MultiDimReductionOp::build(OpBuilder &builder,
492	OperationState &result, Value source,
493	Value acc, ArrayRef<bool> reductionMask,
494	CombiningKind kind) {
495	SmallVector<int64_t> reductionDims;
496	for (const auto &en : llvm::enumerate(reductionMask))
497	if (en.value())
498	reductionDims.push_back(en.index());
499	build(builder, result, kind, source, acc, reductionDims);
500	}
501
502	OpFoldResult MultiDimReductionOp::fold(FoldAdaptor adaptor) {
503	// Single parallel dim, this is a noop.
504	if (getSourceVectorType().getRank() == 1 && !isReducedDim(0))
505	return getSource();
506	return {};
507	}
508
509	std::optional<SmallVector<int64_t, 4>>
510	MultiDimReductionOp::getShapeForUnroll() {
511	return llvm::to_vector<4>(getSourceVectorType().getShape());
512	}
513
514	LogicalResult MultiDimReductionOp::verify() {
515	SmallVector<int64_t> targetShape;
516	SmallVector<bool> scalableDims;
517	Type inferredReturnType;
518	auto sourceScalableDims = getSourceVectorType().getScalableDims();
519	for (auto [dimIdx, dimSize] :
520	llvm::enumerate(getSourceVectorType().getShape()))
521	if (!llvm::any_of(getReductionDims(),
522	[dimIdx = dimIdx](int64_t reductionDimIdx) {
523	return reductionDimIdx == static_cast<int64_t>(dimIdx);
524	})) {
525	targetShape.push_back(dimSize);
526	scalableDims.push_back(sourceScalableDims[dimIdx]);
527	}
528	// TODO: update to also allow 0-d vectors when available.
529	if (targetShape.empty())
530	inferredReturnType = getSourceVectorType().getElementType();
531	else
532	inferredReturnType = VectorType::get(
533	targetShape, getSourceVectorType().getElementType(), scalableDims);
534	if (getType() != inferredReturnType)
535	return emitOpError() << "destination type " << getType()
536	<< " is incompatible with source type "
537	<< getSourceVectorType();
538
539	return success();
540	}
541
542	/// Returns the mask type expected by this operation.
543	Type MultiDimReductionOp::getExpectedMaskType() {
544	auto vecType = getSourceVectorType();
545	return VectorType::get(vecType.getShape(),
546	IntegerType::get(vecType.getContext(), /*width=*/1),
547	vecType.getScalableDims());
548	}
549
550	namespace {
551	// Only unit dimensions that are being reduced are folded. If the dimension is
552	// unit, but not reduced, it is not folded, thereby keeping the output type the
553	// same. If not all dimensions which are reduced are of unit dimension, this
554	// transformation does nothing. This is just a generalization of
555	// ElideSingleElementReduction for ReduceOp.
556	struct ElideUnitDimsInMultiDimReduction
557	: public OpRewritePattern<MultiDimReductionOp> {
558	using OpRewritePattern::OpRewritePattern;
559
560	LogicalResult matchAndRewrite(MultiDimReductionOp reductionOp,
561	PatternRewriter &rewriter) const override {
562	ArrayRef<int64_t> shape = reductionOp.getSourceVectorType().getShape();
563	for (const auto &dim : enumerate(shape)) {
564	if (reductionOp.isReducedDim(dim.index()) && dim.value() != 1)
565	return failure();
566	}
567
568	// Vector mask setup.
569	OpBuilder::InsertionGuard guard(rewriter);
570	Operation *rootOp;
571	Value mask;
572	if (reductionOp.isMasked()) {
573	rewriter.setInsertionPoint(reductionOp.getMaskingOp());
574	rootOp = reductionOp.getMaskingOp();
575	mask = reductionOp.getMaskingOp().getMask();
576	} else {
577	rootOp = reductionOp;
578	}
579
580	Location loc = reductionOp.getLoc();
581	Value acc = reductionOp.getAcc();
582	Value cast;
583	if (auto dstVecType = dyn_cast<VectorType>(reductionOp.getDestType())) {
584	if (mask) {
585	VectorType newMaskType =
586	VectorType::get(dstVecType.getShape(), rewriter.getI1Type(),
587	dstVecType.getScalableDims());
588	mask = rewriter.create<vector::ShapeCastOp>(loc, newMaskType, mask);
589	}
590	cast = rewriter.create<vector::ShapeCastOp>(
591	loc, reductionOp.getDestType(), reductionOp.getSource());
592	} else {
593	// This means we are reducing all the dimensions, and all reduction
594	// dimensions are of size 1. So a simple extraction would do.
595	if (mask)
596	mask = rewriter.create<vector::ExtractOp>(loc, mask);
597	cast = rewriter.create<vector::ExtractOp>(loc, reductionOp.getSource());
598	}
599
600	Value result =
601	vector::makeArithReduction(rewriter, loc, reductionOp.getKind(), acc,
602	cast, /*fastmath=*/nullptr, mask);
603	rewriter.replaceOp(op: rootOp, newValues: result);
604	return success();
605	}
606	};
607	} // namespace
608
609	void MultiDimReductionOp::getCanonicalizationPatterns(
610	RewritePatternSet &results, MLIRContext *context) {
611	results.add<ElideUnitDimsInMultiDimReduction>(context);
612	}
613
614	//===----------------------------------------------------------------------===//
615	// ReductionOp
616	//===----------------------------------------------------------------------===//
617
618	void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
619	CombiningKind kind, Value vector,
620	arith::FastMathFlags fastMathFlags) {
621	build(builder, result, kind, vector, /*acc=*/Value(), fastMathFlags);
622	}
623
624	void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
625	CombiningKind kind, Value vector, Value acc,
626	arith::FastMathFlags fastMathFlags) {
627	build(builder, result,
628	llvm::cast<VectorType>(vector.getType()).getElementType(), kind, vector,
629	acc, fastMathFlags);
630	}
631
632	LogicalResult ReductionOp::verify() {
633	// Verify for 0-D and 1-D vector.
634	int64_t rank = getSourceVectorType().getRank();
635	if (rank > 1)
636	return emitOpError("unsupported reduction rank: ") << rank;
637
638	// Verify supported reduction kind.
639	Type eltType = getDest().getType();
640	if (!isSupportedCombiningKind(getKind(), eltType))
641	return emitOpError("unsupported reduction type '")
642	<< eltType << "' for kind '" << stringifyCombiningKind(getKind())
643	<< "'";
644
645	return success();
646	}
647
648	// MaskableOpInterface methods.
649
650	/// Returns the mask type expected by this operation.
651	Type ReductionOp::getExpectedMaskType() {
652	auto vecType = getSourceVectorType();
653	return VectorType::get(vecType.getShape(),
654	IntegerType::get(vecType.getContext(), /*width=*/1),
655	vecType.getScalableDims());
656	}
657
658	Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
659	OpBuilder &builder, Location loc,
660	Value vector) {
661	switch (op) {
662	case arith::AtomicRMWKind::addf:
663	case arith::AtomicRMWKind::addi:
664	return builder.create<vector::ReductionOp>(vector.getLoc(),
665	CombiningKind::ADD, vector);
666	case arith::AtomicRMWKind::mulf:
667	case arith::AtomicRMWKind::muli:
668	return builder.create<vector::ReductionOp>(vector.getLoc(),
669	CombiningKind::MUL, vector);
670	case arith::AtomicRMWKind::minimumf:
671	return builder.create<vector::ReductionOp>(vector.getLoc(),
672	CombiningKind::MINIMUMF, vector);
673	case arith::AtomicRMWKind::mins:
674	return builder.create<vector::ReductionOp>(vector.getLoc(),
675	CombiningKind::MINSI, vector);
676	case arith::AtomicRMWKind::minu:
677	return builder.create<vector::ReductionOp>(vector.getLoc(),
678	CombiningKind::MINUI, vector);
679	case arith::AtomicRMWKind::maximumf:
680	return builder.create<vector::ReductionOp>(vector.getLoc(),
681	CombiningKind::MAXIMUMF, vector);
682	case arith::AtomicRMWKind::maxs:
683	return builder.create<vector::ReductionOp>(vector.getLoc(),
684	CombiningKind::MAXSI, vector);
685	case arith::AtomicRMWKind::maxu:
686	return builder.create<vector::ReductionOp>(vector.getLoc(),
687	CombiningKind::MAXUI, vector);
688	case arith::AtomicRMWKind::andi:
689	return builder.create<vector::ReductionOp>(vector.getLoc(),
690	CombiningKind::AND, vector);
691	case arith::AtomicRMWKind::ori:
692	return builder.create<vector::ReductionOp>(vector.getLoc(),
693	CombiningKind::OR, vector);
694	// TODO: Add remaining reduction operations.
695	default:
696	(void)emitOptionalError(loc, args: "Reduction operation type not supported");
697	break;
698	}
699	return nullptr;
700	}
701
702	std::optional<SmallVector<int64_t, 4>> ReductionOp::getShapeForUnroll() {
703	return llvm::to_vector<4>(getSourceVectorType().getShape());
704	}
705
706	namespace {
707	struct ElideSingleElementReduction : public OpRewritePattern<ReductionOp> {
708	using OpRewritePattern::OpRewritePattern;
709
710	LogicalResult matchAndRewrite(ReductionOp reductionOp,
711	PatternRewriter &rewriter) const override {
712	// Vector mask setup.
713	OpBuilder::InsertionGuard guard(rewriter);
714	auto maskableOp =
715	cast<vector::MaskableOpInterface>(reductionOp.getOperation());
716	Operation *rootOp;
717	Value mask;
718	if (maskableOp.isMasked()) {
719	rewriter.setInsertionPoint(maskableOp.getMaskingOp());
720	rootOp = maskableOp.getMaskingOp();
721	mask = maskableOp.getMaskingOp().getMask();
722	} else {
723	rootOp = reductionOp;
724	}
725
726	auto vectorType = reductionOp.getSourceVectorType();
727	if (vectorType.getRank() != 0 && vectorType.getDimSize(0) != 1)
728	return failure();
729
730	Location loc = reductionOp.getLoc();
731	if (mask)
732	mask = rewriter.create<ExtractOp>(loc, mask);
733	Value result = rewriter.create<ExtractOp>(loc, reductionOp.getVector());
734
735	if (Value acc = reductionOp.getAcc())
736	result = vector::makeArithReduction(rewriter, loc, reductionOp.getKind(),
737	result, acc,
738	reductionOp.getFastmathAttr(), mask);
739
740	rewriter.replaceOp(op: rootOp, newValues: result);
741	return success();
742	}
743	};
744	} // namespace
745
746	void ReductionOp::getCanonicalizationPatterns(RewritePatternSet &results,
747	MLIRContext *context) {
748	results.add<ElideSingleElementReduction>(context);
749	}
750
751	//===----------------------------------------------------------------------===//
752	// ContractionOp
753	//===----------------------------------------------------------------------===//
754
755	void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
756	Value lhs, Value rhs, Value acc,
757	ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
758	ArrayRef<IteratorType> iteratorTypes) {
759	result.addOperands({lhs, rhs, acc});
760	result.addTypes(acc.getType());
761	result.addAttribute(
762	getIndexingMapsAttrName(result.name),
763	builder.getAffineMapArrayAttr(
764	AffineMap::inferFromExprList(indexingExprs, builder.getContext())));
765	result.addAttribute(
766	getIteratorTypesAttrName(result.name),
767	builder.getArrayAttr(llvm::to_vector(llvm::map_range(
768	iteratorTypes, [&](IteratorType t) -> mlir::Attribute {
769	return IteratorTypeAttr::get(builder.getContext(), t);
770	}))));
771	}
772
773	void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
774	Value lhs, Value rhs, Value acc,
775	ArrayAttr indexingMaps,
776	ArrayAttr iteratorTypes) {
777	build(builder, result, lhs, rhs, acc, indexingMaps, iteratorTypes,
778	ContractionOp::getDefaultKind());
779	}
780
781	void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
782	Value lhs, Value rhs, Value acc,
783	ArrayAttr indexingMaps,
784	ArrayAttr iteratorTypes, CombiningKind kind) {
785	result.addOperands({lhs, rhs, acc});
786	result.addTypes(acc.getType());
787	result.addAttribute(getIndexingMapsAttrName(result.name), indexingMaps);
788	result.addAttribute(getIteratorTypesAttrName(result.name), iteratorTypes);
789	result.addAttribute(getKindAttrName(result.name),
790	CombiningKindAttr::get(builder.getContext(), kind));
791	}
792
793	ParseResult ContractionOp::parse(OpAsmParser &parser, OperationState &result) {
794	OpAsmParser::UnresolvedOperand lhsInfo;
795	OpAsmParser::UnresolvedOperand rhsInfo;
796	OpAsmParser::UnresolvedOperand accInfo;
797	SmallVector<OpAsmParser::UnresolvedOperand, 2> masksInfo;
798	SmallVector<Type, 2> types;
799	Type resultType;
800	auto loc = parser.getCurrentLocation();
801	DictionaryAttr dictAttr;
802	// TODO: Unify linalg op attribute parsing.
803	if (parser.parseAttribute(dictAttr) || parser.parseOperand(lhsInfo) ||
804	parser.parseComma() || parser.parseOperand(rhsInfo) ||
805	parser.parseComma() || parser.parseOperand(accInfo) ||
806	parser.parseTrailingOperandList(masksInfo) ||
807	parser.parseOptionalAttrDict(result.attributes) ||
808	parser.parseColonTypeList(types) ||
809	parser.parseKeywordType("into", resultType) ||
810	parser.resolveOperand(lhsInfo, types[0], result.operands) ||
811	parser.resolveOperand(rhsInfo, types[1], result.operands) ||
812	parser.resolveOperand(accInfo, resultType, result.operands) ||
813	parser.addTypeToList(resultType, result.types))
814	return failure();
815	result.attributes.append(dictAttr.getValue().begin(),
816	dictAttr.getValue().end());
817
818	// Convert array of string into an array of IteratyType enums. This is needed,
819	// because tests still use the old format when 'iterator_types' attribute is
820	// represented as an array of strings.
821	// TODO: Remove this conversion once tests are fixed.
822	auto iteratorTypes = dyn_cast_or_null<ArrayAttr>(
823	result.attributes.get(getIteratorTypesAttrName(result.name)));
824	if (!iteratorTypes) {
825	return parser.emitError(loc)
826	<< "expected " << getIteratorTypesAttrName(result.name)
827	<< " array attribute";
828	}
829
830	SmallVector<Attribute> iteratorTypeAttrs;
831
832	for (StringRef s : iteratorTypes.getAsValueRange<StringAttr>()) {
833	auto maybeIteratorType = symbolizeIteratorType(s);
834	if (!maybeIteratorType.has_value())
835	return parser.emitError(loc) << "unexpected iterator_type (" << s << ")";
836
837	iteratorTypeAttrs.push_back(
838	IteratorTypeAttr::get(parser.getContext(), maybeIteratorType.value()));
839	}
840	result.attributes.set(getIteratorTypesAttrName(result.name),
841	parser.getBuilder().getArrayAttr(iteratorTypeAttrs));
842
843	if (!result.attributes.get(getKindAttrName(result.name))) {
844	result.addAttribute(
845	getKindAttrName(result.name),
846	CombiningKindAttr::get(result.getContext(),
847	ContractionOp::getDefaultKind()));
848	}
849	if (masksInfo.empty())
850	return success();
851	if (masksInfo.size() != 2)
852	return parser.emitError(parser.getNameLoc(),
853	"expected zero or exactly 2 vector mask operands");
854	auto lhsType = llvm::cast<VectorType>(types[0]);
855	auto rhsType = llvm::cast<VectorType>(types[1]);
856	auto maskElementType = parser.getBuilder().getI1Type();
857	std::array<VectorType, 2> maskTypes = {
858	VectorType::Builder(lhsType).setElementType(maskElementType),
859	VectorType::Builder(rhsType).setElementType(maskElementType)};
860	if (parser.resolveOperands(masksInfo, maskTypes, loc, result.operands))
861	return failure();
862	return success();
863	}
864
865	void ContractionOp::print(OpAsmPrinter &p) {
866	// TODO: Unify printing code with linalg ops.
867	auto attrNames = getTraitAttrNames();
868	llvm::StringSet<> traitAttrsSet;
869	traitAttrsSet.insert_range(attrNames);
870	SmallVector<NamedAttribute, 8> attrs;
871	for (auto attr : (*this)->getAttrs()) {
872	if (attr.getName() == getIteratorTypesAttrName()) {
873	auto iteratorTypes =
874	llvm::cast<ArrayAttr>(attr.getValue())
875	.getAsValueRange<IteratorTypeAttr, IteratorType>();
876	// Convert IteratorType enums into the string representation. This is
877	// needed, because tests still use the old format when 'iterator_types'
878	// attribute is represented as an array of strings.
879	// TODO: Remove this conversion once tests are fixed.
880	SmallVector<Attribute> iteratorTypeNames = llvm::to_vector(
881	llvm::map_range(iteratorTypes, [&](IteratorType t) -> Attribute {
882	return StringAttr::get(getContext(), stringifyIteratorType(t));
883	}));
884
885	attrs.emplace_back(getIteratorTypesAttrName(),
886	ArrayAttr::get(getContext(), iteratorTypeNames));
887	} else if (traitAttrsSet.count(attr.getName().strref()) > 0)
888	attrs.push_back(attr);
889	}
890
891	auto dictAttr = DictionaryAttr::get(getContext(), attrs);
892	p << " " << dictAttr << " " << getLhs() << ", ";
893	p << getRhs() << ", " << getAcc();
894
895	p.printOptionalAttrDict((*this)->getAttrs(), attrNames);
896	p << " : " << getLhs().getType() << ", " << getRhs().getType() << " into "
897	<< getResultType();
898	}
899
900	static bool verifyDimMap(VectorType lhsType, VectorType rhsType,
901	const std::vector<std::pair<int64_t, int64_t>> &map) {
902	for (auto &dimPair : map) {
903	if (dimPair.first < 0 || dimPair.first >= lhsType.getRank() ||
904	dimPair.second < 0 || dimPair.second >= rhsType.getRank() ||
905	lhsType.getDimSize(dimPair.first) != rhsType.getDimSize(dimPair.second))
906	return false;
907	}
908	return true;
909	}
910
911	static LogicalResult verifyOutputShape(
912	ContractionOp op, VectorType lhsType, VectorType rhsType, Type accType,
913	Type resType,
914	const std::vector<std::pair<int64_t, int64_t>> &contractingDimMap,
915	const std::vector<std::pair<int64_t, int64_t>> &batchDimMap) {
916	DenseSet<int64_t> lhsContractingDimSet;
917	DenseSet<int64_t> rhsContractingDimSet;
918	for (auto &dimPair : contractingDimMap) {
919	lhsContractingDimSet.insert(V: dimPair.first);
920	rhsContractingDimSet.insert(V: dimPair.second);
921	}
922	DenseSet<int64_t> rhsBatchDimSet(llvm::from_range,
923	llvm::make_second_range(c: batchDimMap));
924
925	// Add free and batch dimensions from 'lhsType' to 'expectedResultDims'.
926	SmallVector<int64_t, 4> expectedResultDims;
927	for (int64_t i = 0, e = lhsType.getRank(); i < e; ++i) {
928	if (lhsContractingDimSet.count(V: i) > 0)
929	continue;
930	expectedResultDims.push_back(Elt: lhsType.getDimSize(i));
931	}
932
933	// Add free dimensions from 'rhsType' to 'expectedResultDims'.
934	for (int64_t i = 0, e = rhsType.getRank(); i < e; ++i) {
935	if (rhsContractingDimSet.count(V: i) > 0 || rhsBatchDimSet.count(V: i) > 0)
936	continue;
937	expectedResultDims.push_back(Elt: rhsType.getDimSize(i));
938	}
939
940	// Verify 'expectedResultDims'.
941	if (expectedResultDims.empty()) {
942	// No batch or free dimension implies a scalar result.
943	if (llvm::isa<VectorType>(Val: resType) || llvm::isa<VectorType>(Val: accType))
944	return op.emitOpError("invalid accumulator/result vector shape");
945	} else {
946	// At least one batch or free dimension implies a vector result.
947	auto resVectorType = llvm::dyn_cast<VectorType>(resType);
948	auto accVectorType = llvm::dyn_cast<VectorType>(accType);
949	if (!resVectorType || !accVectorType)
950	return op.emitOpError("invalid accumulator/result vector shape");
951
952	// Infer expected result vector type. Lhs + rhs map and lhs + rhs vector
953	// types fully define the result vector type. This assumes the affine maps
954	// are well-formed, which must have been verified already.
955	MLIRContext *ctx = op.getContext();
956	AffineMap lhsMap = op.getIndexingMapsArray()[0];
957	AffineMap rhsMap = op.getIndexingMapsArray()[1];
958	if (getUnusedDimsBitVector(maps: {lhsMap, rhsMap}).any())
959	return op.emitOpError(
960	"expected all dimensions to be either a LHS or a RHS dimension");
961	SmallVector<AffineExpr, 4> extents(lhsMap.getNumInputs());
962	for (auto pair :
963	{std::make_pair(lhsType, lhsMap), std::make_pair(rhsType, rhsMap)}) {
964	VectorType v = pair.first;
965	auto map = pair.second;
966	for (unsigned idx = 0, e = v.getRank(); idx < e; ++idx) {
967	unsigned pos = map.getDimPosition(idx);
968	if (!extents[pos])
969	extents[pos] = getAffineConstantExpr(v.getShape()[idx], ctx);
970	}
971	}
972	if (!llvm::all_of(Range&: extents, P: [](AffineExpr e) { return e; }))
973	return op.emitOpError("expected all dimensions to get an extent as "
974	"either a LHS or a RHS dimension");
975
976	AffineMap resMap = op.getIndexingMapsArray()[2];
977	auto extentsMap = AffineMap::get(/*dimCount=*/extents.size(),
978	/*symbolCount=*/0, results: extents, context: ctx);
979	// Compose the resMap with the extentsMap, which is a constant map.
980	AffineMap expectedMap = simplifyAffineMap(resMap.compose(extentsMap));
981	assert(llvm::all_of(expectedMap.getResults(),
982	llvm::IsaPred<AffineConstantExpr>) &&
983	"expected constant extent along all dimensions.");
984	// Extract the expected shape and build the type.
985	auto expectedShape = llvm::to_vector<4>(
986	Range: llvm::map_range(C: expectedMap.getResults(), F: [](AffineExpr e) {
987	return cast<AffineConstantExpr>(Val&: e).getValue();
988	}));
989	auto expected =
990	VectorType::get(expectedShape, resVectorType.getElementType(),
991	resVectorType.getScalableDims());
992	if (resVectorType != expected || accVectorType != expected)
993	return op.emitOpError(
994	"invalid accumulator/result vector shape, expected: ")
995	<< expected;
996	}
997	return success();
998	}
999
1000	LogicalResult ContractionOp::verify() {
1001	VectorType lhsType = getLhsType();
1002	VectorType rhsType = getRhsType();
1003	Type accType = getAccType();
1004	Type resType = getResultType();
1005
1006	if (llvm::isa<IntegerType>(lhsType.getElementType())) {
1007	if (!lhsType.getElementType().isSignlessInteger())
1008	return emitOpError("only supports signless integer types");
1009	}
1010
1011	// Verify that an indexing map was specified for each vector operand.
1012	if (getIndexingMapsArray().size() != 3)
1013	return emitOpError("expected an indexing map for each vector operand");
1014
1015	// Verify that each index map has 'numIterators' inputs, no symbols, and
1016	// that the number of map outputs equals the rank of its associated
1017	// vector operand.
1018	unsigned numIterators = getIteratorTypes().getValue().size();
1019	for (const auto &it : llvm::enumerate(getIndexingMapsArray())) {
1020	auto index = it.index();
1021	auto map = it.value();
1022	if (map.getNumSymbols() != 0)
1023	return emitOpError("expected indexing map ")
1024	<< index << " to have no symbols";
1025	auto vectorType = llvm::dyn_cast<VectorType>(getOperand(index).getType());
1026	unsigned rank = vectorType ? vectorType.getShape().size() : 0;
1027	// Verify that the map has the right number of inputs, outputs, and indices.
1028	// This also correctly accounts for (..) -> () for rank-0 results.
1029	if (map.getNumDims() != numIterators)
1030	return emitOpError("expected indexing map ")
1031	<< index << " to have " << numIterators << " number of inputs";
1032	if (map.getNumResults() != rank)
1033	return emitOpError("expected indexing map ")
1034	<< index << " to have " << rank << " number of outputs";
1035	if (!map.isProjectedPermutation())
1036	return emitOpError("expected indexing map ")
1037	<< index << " to be a projected permutation of its inputs";
1038	}
1039
1040	auto contractingDimMap = getContractingDimMap();
1041	auto batchDimMap = getBatchDimMap();
1042
1043	// Verify at least one contracting dimension pair was specified.
1044	if (contractingDimMap.empty())
1045	return emitOpError("expected at least one contracting dimension pair");
1046
1047	// Verify contracting dimension map was properly constructed.
1048	if (!verifyDimMap(lhsType, rhsType, contractingDimMap))
1049	return emitOpError("invalid contracting dimension map");
1050
1051	// Verify batch dimension map was properly constructed.
1052	if (!verifyDimMap(lhsType, rhsType, batchDimMap))
1053	return emitOpError("invalid batch dimension map");
1054
1055	// Verify 'accType' and 'resType' shape.
1056	if (failed(verifyOutputShape(*this, lhsType, rhsType, accType, resType,
1057	contractingDimMap, batchDimMap)))
1058	return failure();
1059
1060	// Verify supported combining kind.
1061	auto vectorType = llvm::dyn_cast<VectorType>(resType);
1062	auto elementType = vectorType ? vectorType.getElementType() : resType;
1063	if (!isSupportedCombiningKind(getKind(), elementType))
1064	return emitOpError("unsupported contraction type");
1065
1066	return success();
1067	}
1068
1069	// MaskableOpInterface methods.
1070
1071	/// Returns the mask type expected by this operation. Mostly used for
1072	/// verification purposes. It requires the operation to be vectorized."
1073	Type ContractionOp::getExpectedMaskType() {
1074	auto indexingMaps = this->getIndexingMapsArray();
1075	AffineMap lhsIdxMap = indexingMaps[0];
1076	AffineMap rhsIdxMap = indexingMaps[1];
1077	VectorType lhsType = this->getLhsType();
1078	VectorType rhsType = this->getRhsType();
1079
1080	unsigned numVecDims = lhsIdxMap.getNumDims();
1081	SmallVector<int64_t> maskShape(numVecDims, ShapedType::kDynamic);
1082	SmallVector<bool> maskShapeScalableDims(numVecDims, false);
1083
1084	// Using the information in the indexing maps, extract the size of each
1085	// dimension in the vector.contract operation from the two input operands.
1086	for (auto [dimIdx, dimSize] : llvm::enumerate(lhsType.getShape())) {
1087	maskShape[lhsIdxMap.getDimPosition(dimIdx)] = dimSize;
1088	maskShapeScalableDims[lhsIdxMap.getDimPosition(dimIdx)] =
1089	lhsType.getScalableDims()[dimIdx];
1090	}
1091	for (auto [dimIdx, dimSize] : llvm::enumerate(rhsType.getShape())) {
1092	maskShape[rhsIdxMap.getDimPosition(dimIdx)] = dimSize;
1093	maskShapeScalableDims[rhsIdxMap.getDimPosition(dimIdx)] =
1094	rhsType.getScalableDims()[dimIdx];
1095	}
1096
1097	assert(!ShapedType::isDynamicShape(maskShape) &&
1098	"Mask shape couldn't be computed");
1099
1100	return VectorType::get(maskShape,
1101	IntegerType::get(lhsType.getContext(), /*width=*/1),
1102	maskShapeScalableDims);
1103	}
1104
1105	SmallVector<StringRef> ContractionOp::getTraitAttrNames() {
1106	return SmallVector<StringRef>{getIndexingMapsAttrName(),
1107	getIteratorTypesAttrName(), getKindAttrName()};
1108	}
1109
1110	static int64_t getResultIndex(AffineMap map, AffineExpr targetExpr) {
1111	for (int64_t i = 0, e = map.getNumResults(); i < e; ++i)
1112	if (targetExpr == map.getResult(idx: i))
1113	return i;
1114	return -1;
1115	}
1116
1117	static std::vector<std::pair<int64_t, int64_t>>
1118	getDimMap(ArrayRef<AffineMap> indexingMaps, ArrayAttr iteratorTypes,
1119	IteratorType targetIteratorType, MLIRContext *context) {
1120	std::vector<std::pair<int64_t, int64_t>> dimMap;
1121	for (const auto &it : llvm::enumerate(iteratorTypes)) {
1122	auto iteratorType = llvm::cast<IteratorTypeAttr>(it.value()).getValue();
1123	if (iteratorType != targetIteratorType)
1124	continue;
1125	// Search lhs/rhs map results for 'targetExpr'.
1126	auto targetExpr = getAffineDimExpr(it.index(), context);
1127	int64_t lhsDim = getResultIndex(indexingMaps[0], targetExpr);
1128	int64_t rhsDim = getResultIndex(indexingMaps[1], targetExpr);
1129	if (lhsDim >= 0 && rhsDim >= 0)
1130	dimMap.emplace_back(lhsDim, rhsDim);
1131	}
1132	return dimMap;
1133	}
1134
1135	void ContractionOp::getIterationBounds(
1136	SmallVectorImpl<int64_t> &iterationBounds) {
1137	auto lhsShape = getLhsType().getShape();
1138	auto resVectorType = llvm::dyn_cast<VectorType>(getResultType());
1139	SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
1140	for (const auto &it : llvm::enumerate(getIteratorTypes())) {
1141	// Search lhs/rhs map results for 'targetExpr'.
1142	auto targetExpr = getAffineDimExpr(it.index(), getContext());
1143	auto iteratorType = llvm::cast<IteratorTypeAttr>(it.value()).getValue();
1144	if (iteratorType == IteratorType::reduction) {
1145	// Get reduction dim size from lhs shape (same size in rhsShape).
1146	int64_t lhsDimIndex = getResultIndex(indexingMaps[0], targetExpr);
1147	assert(lhsDimIndex >= 0);
1148	iterationBounds.push_back(lhsShape[lhsDimIndex]);
1149	continue;
1150	}
1151	// Get parallel dimension size from result shape.
1152	int64_t resDimIndex = getResultIndex(indexingMaps[2], targetExpr);
1153	assert(resDimIndex >= 0);
1154	assert(resVectorType != nullptr);
1155	iterationBounds.push_back(resVectorType.getShape()[resDimIndex]);
1156	}
1157	}
1158
1159	void ContractionOp::getIterationIndexMap(
1160	std::vector<DenseMap<int64_t, int64_t>> &iterationIndexMap) {
1161	unsigned numMaps = getIndexingMapsArray().size();
1162	iterationIndexMap.resize(numMaps);
1163	for (const auto &it : llvm::enumerate(getIndexingMapsArray())) {
1164	auto index = it.index();
1165	auto map = it.value();
1166	for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
1167	auto dim = cast<AffineDimExpr>(map.getResult(i));
1168	iterationIndexMap[index][dim.getPosition()] = i;
1169	}
1170	}
1171	}
1172
1173	std::vector<std::pair<int64_t, int64_t>> ContractionOp::getContractingDimMap() {
1174	SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
1175	return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::reduction,
1176	getContext());
1177	}
1178
1179	std::vector<std::pair<int64_t, int64_t>> ContractionOp::getBatchDimMap() {
1180	SmallVector<AffineMap, 4> indexingMaps(getIndexingMapsArray());
1181	return getDimMap(indexingMaps, getIteratorTypes(), IteratorType::parallel,
1182	getContext());
1183	}
1184
1185	std::optional<SmallVector<int64_t, 4>> ContractionOp::getShapeForUnroll() {
1186	SmallVector<int64_t, 4> shape;
1187	getIterationBounds(shape);
1188	return shape;
1189	}
1190
1191	/// Return a fused vector::ContractionOp which represents a patterns such as:
1192	///
1193	/// ```mlir
1194	/// %c0 = vector.constant 0: ...
1195	/// %c = vector.contract %a, %b, %c0: ...
1196	/// %e = add %c, %d: ...
1197	/// ```
1198	///
1199	/// by:
1200	///
1201	/// ```mlir
1202	/// %e = vector.contract %a, %b, %d: ...
1203	/// ```
1204	///
1205	/// Return null if the canonicalization does not apply.
1206	// TODO: This should be a folding of Add into Contract in core but while they
1207	// live in different dialects, it is not possible without unnatural
1208	// dependencies.
1209	template <typename AddOpType>
1210	struct CanonicalizeContractAdd : public OpRewritePattern<AddOpType> {
1211	using OpRewritePattern<AddOpType>::OpRewritePattern;
1212
1213	LogicalResult matchAndRewrite(AddOpType addOp,
1214	PatternRewriter &rewriter) const override {
1215	auto canonicalize = [&](Value maybeContraction,
1216	Value otherOperand) -> vector::ContractionOp {
1217	vector::ContractionOp contractionOp =
1218	dyn_cast_or_null<vector::ContractionOp>(
1219	maybeContraction.getDefiningOp());
1220	if (!contractionOp)
1221	return vector::ContractionOp();
1222	if (auto maybeZero = dyn_cast_or_null<arith::ConstantOp>(
1223	contractionOp.getAcc().getDefiningOp())) {
1224	if (maybeZero.getValue() ==
1225	rewriter.getZeroAttr(type: contractionOp.getAcc().getType())) {
1226	IRMapping bvm;
1227	bvm.map(contractionOp.getAcc(), otherOperand);
1228	auto newContraction =
1229	cast<vector::ContractionOp>(rewriter.clone(*contractionOp, bvm));
1230	rewriter.replaceOp(addOp, newContraction.getResult());
1231	return newContraction;
1232	}
1233	}
1234	return vector::ContractionOp();
1235	};
1236
1237	Value a = addOp->getOperand(0), b = addOp->getOperand(1);
1238	vector::ContractionOp contract = canonicalize(a, b);
1239	contract = contract ? contract : canonicalize(b, a);
1240	return contract ? success() : failure();
1241	}
1242	};
1243
1244	void ContractionOp::getCanonicalizationPatterns(RewritePatternSet &results,
1245	MLIRContext *context) {
1246	results.add<CanonicalizeContractAdd<arith::AddIOp>,
1247	CanonicalizeContractAdd<arith::AddFOp>>(context);
1248	}
1249
1250	//===----------------------------------------------------------------------===//
1251	// ExtractElementOp
1252	//===----------------------------------------------------------------------===//
1253
1254	void ExtractElementOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
1255	SetIntRangeFn setResultRanges) {
1256	setResultRanges(getResult(), argRanges.front());
1257	}
1258
1259	void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
1260	Value source) {
1261	result.addOperands({source});
1262	result.addTypes(llvm::cast<VectorType>(source.getType()).getElementType());
1263	}
1264
1265	LogicalResult vector::ExtractElementOp::verify() {
1266	VectorType vectorType = getSourceVectorType();
1267	if (vectorType.getRank() == 0) {
1268	if (getPosition())
1269	return emitOpError("expected position to be empty with 0-D vector");
1270	return success();
1271	}
1272	if (vectorType.getRank() != 1)
1273	return emitOpError("unexpected >1 vector rank");
1274	if (!getPosition())
1275	return emitOpError("expected position for 1-D vector");
1276	return success();
1277	}
1278
1279	OpFoldResult vector::ExtractElementOp::fold(FoldAdaptor adaptor) {
1280	// Skip the 0-D vector here now.
1281	if (!adaptor.getPosition())
1282	return {};
1283
1284	// Fold extractelement (splat X) -> X.
1285	if (auto splat = getVector().getDefiningOp<vector::SplatOp>())
1286	return splat.getInput();
1287
1288	// Fold extractelement(broadcast(X)) -> X.
1289	if (auto broadcast = getVector().getDefiningOp<vector::BroadcastOp>())
1290	if (!llvm::isa<VectorType>(broadcast.getSource().getType()))
1291	return broadcast.getSource();
1292
1293	auto src = dyn_cast_or_null<DenseElementsAttr>(adaptor.getVector());
1294	auto pos = dyn_cast_or_null<IntegerAttr>(adaptor.getPosition());
1295	if (!pos || !src)
1296	return {};
1297
1298	auto srcElements = src.getValues<Attribute>();
1299
1300	uint64_t posIdx = pos.getInt();
1301	if (posIdx >= srcElements.size())
1302	return {};
1303
1304	return srcElements[posIdx];
1305	}
1306
1307	// Returns `true` if `index` is either within [0, maxIndex) or equal to
1308	// `poisonValue`.
1309	static bool isValidPositiveIndexOrPoison(int64_t index, int64_t poisonValue,
1310	int64_t maxIndex) {
1311	return index == poisonValue || (index >= 0 && index < maxIndex);
1312	}
1313
1314	//===----------------------------------------------------------------------===//
1315	// ExtractOp
1316	//===----------------------------------------------------------------------===//
1317
1318	void ExtractOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
1319	SetIntRangeFn setResultRanges) {
1320	setResultRanges(getResult(), argRanges.front());
1321	}
1322
1323	void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1324	Value source) {
1325	auto vectorTy = cast<VectorType>(source.getType());
1326	build(builder, result, source, SmallVector<int64_t>(vectorTy.getRank(), 0));
1327	}
1328
1329	void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1330	Value source, int64_t position) {
1331	build(builder, result, source, ArrayRef<int64_t>{position});
1332	}
1333
1334	void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1335	Value source, OpFoldResult position) {
1336	build(builder, result, source, ArrayRef<OpFoldResult>{position});
1337	}
1338
1339	void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1340	Value source, ArrayRef<int64_t> position) {
1341	build(builder, result, source, /*dynamic_position=*/ArrayRef<Value>(),
1342	builder.getDenseI64ArrayAttr(position));
1343	}
1344
1345	void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
1346	Value source, ArrayRef<OpFoldResult> position) {
1347	SmallVector<int64_t> staticPos;
1348	SmallVector<Value> dynamicPos;
1349	dispatchIndexOpFoldResults(position, dynamicPos, staticPos);
1350	build(builder, result, source, dynamicPos,
1351	builder.getDenseI64ArrayAttr(staticPos));
1352	}
1353
1354	LogicalResult
1355	ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
1356	ExtractOp::Adaptor adaptor,
1357	SmallVectorImpl<Type> &inferredReturnTypes) {
1358	auto vectorType = llvm::cast<VectorType>(adaptor.getVector().getType());
1359	if (static_cast<int64_t>(adaptor.getStaticPosition().size()) ==
1360	vectorType.getRank()) {
1361	inferredReturnTypes.push_back(vectorType.getElementType());
1362	} else {
1363	auto n = std::min<size_t>(adaptor.getStaticPosition().size(),
1364	vectorType.getRank());
1365	inferredReturnTypes.push_back(VectorType::get(
1366	vectorType.getShape().drop_front(n), vectorType.getElementType(),
1367	vectorType.getScalableDims().drop_front(n)));
1368	}
1369	return success();
1370	}
1371
1372	bool ExtractOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
1373	// Allow extracting 1-element vectors instead of scalars.
1374	auto isCompatible = [](TypeRange l, TypeRange r) {
1375	auto vectorType = llvm::dyn_cast<VectorType>(l.front());
1376	return vectorType && vectorType.getShape().equals({1}) &&
1377	vectorType.getElementType() == r.front();
1378	};
1379	if (l.size() == 1 && r.size() == 1 &&
1380	(isCompatible(l, r) || isCompatible(r, l)))
1381	return true;
1382	return l == r;
1383	}
1384
1385	LogicalResult vector::ExtractOp::verify() {
1386	// Note: This check must come before getMixedPosition() to prevent a crash.
1387	auto dynamicMarkersCount =
1388	llvm::count_if(getStaticPosition(), ShapedType::isDynamic);
1389	if (static_cast<size_t>(dynamicMarkersCount) != getDynamicPosition().size())
1390	return emitOpError(
1391	"mismatch between dynamic and static positions (kDynamic marker but no "
1392	"corresponding dynamic position) -- this can only happen due to an "
1393	"incorrect fold/rewrite");
1394	auto position = getMixedPosition();
1395	if (position.size() > static_cast<unsigned>(getSourceVectorType().getRank()))
1396	return emitOpError(
1397	"expected position attribute of rank no greater than vector rank");
1398	for (auto [idx, pos] : llvm::enumerate(position)) {
1399	if (auto attr = dyn_cast<Attribute>(pos)) {
1400	int64_t constIdx = cast<IntegerAttr>(attr).getInt();
1401	if (!isValidPositiveIndexOrPoison(
1402	constIdx, kPoisonIndex, getSourceVectorType().getDimSize(idx))) {
1403	return emitOpError("expected position attribute #")
1404	<< (idx + 1)
1405	<< " to be a non-negative integer smaller than the "
1406	"corresponding vector dimension or poison (-1)";
1407	}
1408	}
1409	}
1410	return success();
1411	}
1412
1413	template <typename IntType>
1414	static SmallVector<IntType> extractVector(ArrayAttr arrayAttr) {
1415	return llvm::to_vector<4>(llvm::map_range(
1416	arrayAttr.getAsRange<IntegerAttr>(),
1417	[](IntegerAttr attr) { return static_cast<IntType>(attr.getInt()); }));
1418	}
1419
1420	/// Fold the result of chains of ExtractOp in place by simply concatenating the
1421	/// positions.
1422	static LogicalResult foldExtractOpFromExtractChain(ExtractOp extractOp) {
1423	if (!extractOp.getVector().getDefiningOp<ExtractOp>())
1424	return failure();
1425
1426	// TODO: Canonicalization for dynamic position not implemented yet.
1427	if (extractOp.hasDynamicPosition())
1428	return failure();
1429
1430	SmallVector<int64_t> globalPosition;
1431	ExtractOp currentOp = extractOp;
1432	ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
1433	globalPosition.append(in_start: extrPos.rbegin(), in_end: extrPos.rend());
1434	while (ExtractOp nextOp = currentOp.getVector().getDefiningOp<ExtractOp>()) {
1435	currentOp = nextOp;
1436	// TODO: Canonicalization for dynamic position not implemented yet.
1437	if (currentOp.hasDynamicPosition())
1438	return failure();
1439	ArrayRef<int64_t> extrPos = currentOp.getStaticPosition();
1440	globalPosition.append(in_start: extrPos.rbegin(), in_end: extrPos.rend());
1441	}
1442	extractOp.setOperand(0, currentOp.getVector());
1443	// OpBuilder is only used as a helper to build an I64ArrayAttr.
1444	OpBuilder b(extractOp.getContext());
1445	std::reverse(first: globalPosition.begin(), last: globalPosition.end());
1446	extractOp.setStaticPosition(globalPosition);
1447	return success();
1448	}
1449
1450	namespace {
1451	/// Fold an ExtractOp that is fed by a chain of InsertOps and TransposeOps.
1452	/// Walk back a chain of InsertOp/TransposeOp until we hit a match.
1453	/// Compose TransposeOp permutations as we walk back.
1454	/// This helper class keeps an updated extraction position `extractPosition`
1455	/// with extra trailing sentinels.
1456	/// The sentinels encode the internal transposition status of the result vector.
1457	/// As we iterate, extractPosition is permuted and updated.
1458	class ExtractFromInsertTransposeChainState {
1459	public:
1460	ExtractFromInsertTransposeChainState(ExtractOp e);
1461
1462	/// Iterate over producing insert and transpose ops until we find a fold.
1463	Value fold();
1464
1465	private:
1466	/// Return true if the vector at position `a` is contained within the vector
1467	/// at position `b`. Under insert/extract semantics, this is the same as `a`
1468	/// is a prefix of `b`.
1469	template <typename ContainerA, typename ContainerB>
1470	bool isContainedWithin(const ContainerA &a, const ContainerB &b) {
1471	return a.size() <= b.size() &&
1472	std::equal(a.begin(), a.begin() + a.size(), b.begin());
1473	}
1474
1475	/// Return true if the vector at position `a` intersects the vector at
1476	/// position `b`. Under insert/extract semantics, this is the same as equality
1477	/// of all entries of `a` that are >=0 with the corresponding entries of b.
1478	/// Comparison is on the common prefix (i.e. zip).
1479	template <typename ContainerA, typename ContainerB>
1480	bool intersectsWhereNonNegative(const ContainerA &a, const ContainerB &b) {
1481	for (auto [elemA, elemB] : llvm::zip(a, b)) {
1482	if (elemA < 0 || elemB < 0)
1483	continue;
1484	if (elemA != elemB)
1485	return false;
1486	}
1487	return true;
1488	}
1489
1490	/// Folding is only possible in the absence of an internal permutation in the
1491	/// result vector.
1492	bool canFold() {
1493	return (sentinels == ArrayRef(extractPosition).drop_front(N: extractedRank));
1494	}
1495
1496	// Helper to get the next defining op of interest.
1497	void updateStateForNextIteration(Value v) {
1498	nextInsertOp = v.getDefiningOp<vector::InsertOp>();
1499	nextTransposeOp = v.getDefiningOp<vector::TransposeOp>();
1500	};
1501
1502	// Case 1. If we hit a transpose, just compose the map and iterate.
1503	// Invariant: insert + transpose do not change rank, we can always compose.
1504	LogicalResult handleTransposeOp();
1505
1506	// Case 2: the insert position matches extractPosition exactly, early return.
1507	LogicalResult handleInsertOpWithMatchingPos(Value &res);
1508
1509	/// Case 3: if the insert position is a prefix of extractPosition, extract a
1510	/// portion of the source of the insert.
1511	/// Example:
1512	/// ```
1513	/// %ins = vector.insert %source, %vest[1]: vector<3x4> into vector<2x3x4x5>
1514	/// // extractPosition == [1, 2, 3]
1515	/// %ext = vector.extract %ins[1, 0]: vector<5> from vector<3x4x5>
1516	/// // can fold to vector.extract %source[0, 3]
1517	/// %ext = vector.extract %source[3]: vector<6> from vector<5x6>
1518	/// ```
1519	/// To traverse through %source, we need to set the leading dims to 0 and
1520	/// drop the extra leading dims.
1521	/// This method updates the internal state.
1522	LogicalResult handleInsertOpWithPrefixPos(Value &res);
1523
1524	/// Try to fold in place to extract(source, extractPosition) and return the
1525	/// folded result. Return null if folding is not possible (e.g. due to an
1526	/// internal transposition in the result).
1527	Value tryToFoldExtractOpInPlace(Value source);
1528
1529	ExtractOp extractOp;
1530	int64_t vectorRank;
1531	int64_t extractedRank;
1532
1533	InsertOp nextInsertOp;
1534	TransposeOp nextTransposeOp;
1535
1536	/// Sentinel values that encode the internal permutation status of the result.
1537	/// They are set to (-1, ... , -k) at the beginning and appended to
1538	/// `extractPosition`.
1539	/// In the end, the tail of `extractPosition` must be exactly `sentinels` to
1540	/// ensure that there is no internal transposition.
1541	/// Internal transposition cannot be accounted for with a folding pattern.
1542	// TODO: We could relax the internal transposition with an extra transposition
1543	// operation in a future canonicalizer.
1544	SmallVector<int64_t> sentinels;
1545	SmallVector<int64_t> extractPosition;
1546	};
1547	} // namespace
1548
1549	ExtractFromInsertTransposeChainState::ExtractFromInsertTransposeChainState(
1550	ExtractOp e)
1551	: extractOp(e), vectorRank(extractOp.getSourceVectorType().getRank()),
1552	extractedRank(extractOp.getNumIndices()) {
1553	assert(vectorRank >= extractedRank && "Extracted position overflow");
1554	sentinels.reserve(N: vectorRank - extractedRank);
1555	for (int64_t i = 0, e = vectorRank - extractedRank; i < e; ++i)
1556	sentinels.push_back(Elt: -(i + 1));
1557	extractPosition.assign(extractOp.getStaticPosition().begin(),
1558	extractOp.getStaticPosition().end());
1559	llvm::append_range(C&: extractPosition, R&: sentinels);
1560	}
1561
1562	// Case 1. If we hit a transpose, just compose the map and iterate.
1563	// Invariant: insert + transpose do not change rank, we can always compose.
1564	LogicalResult ExtractFromInsertTransposeChainState::handleTransposeOp() {
1565	// TODO: Canonicalization for dynamic position not implemented yet.
1566	if (extractOp.hasDynamicPosition())
1567	return failure();
1568
1569	if (!nextTransposeOp)
1570	return failure();
1571	AffineMap m = inversePermutation(AffineMap::getPermutationMap(
1572	nextTransposeOp.getPermutation(), extractOp.getContext()));
1573	extractPosition = applyPermutationMap(map: m, source: ArrayRef(extractPosition));
1574	return success();
1575	}
1576
1577	// Case 2: the insert position matches extractPosition exactly, early return.
1578	LogicalResult
1579	ExtractFromInsertTransposeChainState::handleInsertOpWithMatchingPos(
1580	Value &res) {
1581	// TODO: Canonicalization for dynamic position not implemented yet.
1582	if (extractOp.hasDynamicPosition() || nextInsertOp.hasDynamicPosition())
1583	return failure();
1584
1585	ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
1586	if (insertedPos != llvm::ArrayRef(extractPosition).take_front(N: extractedRank))
1587	return failure();
1588	// Case 2.a. early-exit fold.
1589	res = nextInsertOp.getValueToStore();
1590	// Case 2.b. if internal transposition is present, canFold will be false.
1591	return success(IsSuccess: canFold());
1592	}
1593
1594	/// Case 3: if inserted position is a prefix of extractPosition,
1595	/// extract a portion of the source of the insertion.
1596	/// This method updates the internal state.
1597	LogicalResult
1598	ExtractFromInsertTransposeChainState::handleInsertOpWithPrefixPos(Value &res) {
1599	// TODO: Canonicalization for dynamic position not implemented yet.
1600	if (extractOp.hasDynamicPosition() || nextInsertOp.hasDynamicPosition())
1601	return failure();
1602
1603	ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
1604	if (!isContainedWithin(a: insertedPos, b: extractPosition))
1605	return failure();
1606	// Set leading dims to zero.
1607	std::fill_n(first: extractPosition.begin(), n: insertedPos.size(), value: 0);
1608	// Drop extra leading dims.
1609	extractPosition.erase(CS: extractPosition.begin(),
1610	CE: extractPosition.begin() + insertedPos.size());
1611	extractedRank = extractPosition.size() - sentinels.size();
1612	// Case 3.a. early-exit fold (break and delegate to post-while path).
1613	res = nextInsertOp.getValueToStore();
1614	// Case 3.b. if internal transposition is present, canFold will be false.
1615	return success();
1616	}
1617
1618	/// Try to fold in place to extract(source, extractPosition) and return the
1619	/// folded result. Return null if folding is not possible (e.g. due to an
1620	/// internal transposition in the result).
1621	Value ExtractFromInsertTransposeChainState::tryToFoldExtractOpInPlace(
1622	Value source) {
1623	// TODO: Canonicalization for dynamic position not implemented yet.
1624	if (extractOp.hasDynamicPosition())
1625	return Value();
1626
1627	// If we can't fold (either internal transposition, or nothing to fold), bail.
1628	bool nothingToFold = (source == extractOp.getVector());
1629	if (nothingToFold || !canFold())
1630	return Value();
1631
1632	// Otherwise, fold by updating the op inplace and return its result.
1633	OpBuilder b(extractOp.getContext());
1634	extractOp.setStaticPosition(
1635	ArrayRef(extractPosition).take_front(extractedRank));
1636	extractOp.getVectorMutable().assign(source);
1637	return extractOp.getResult();
1638	}
1639
1640	/// Iterate over producing insert and transpose ops until we find a fold.
1641	Value ExtractFromInsertTransposeChainState::fold() {
1642	// TODO: Canonicalization for dynamic position not implemented yet.
1643	if (extractOp.hasDynamicPosition())
1644	return Value();
1645
1646	Value valueToExtractFrom = extractOp.getVector();
1647	updateStateForNextIteration(v: valueToExtractFrom);
1648	while (nextInsertOp || nextTransposeOp) {
1649	// Case 1. If we hit a transpose, just compose the map and iterate.
1650	// Invariant: insert + transpose do not change rank, we can always compose.
1651	if (succeeded(Result: handleTransposeOp())) {
1652	valueToExtractFrom = nextTransposeOp.getVector();
1653	updateStateForNextIteration(v: valueToExtractFrom);
1654	continue;
1655	}
1656
1657	Value result;
1658	// Case 2: the position match exactly.
1659	if (succeeded(Result: handleInsertOpWithMatchingPos(res&: result)))
1660	return result;
1661
1662	// Case 3: if the inserted position is a prefix of extractPosition, we can
1663	// just extract a portion of the source of the insert.
1664	if (succeeded(Result: handleInsertOpWithPrefixPos(res&: result)))
1665	return tryToFoldExtractOpInPlace(source: result);
1666
1667	// Case 4: extractPositionRef intersects insertedPosRef on non-sentinel
1668	// values. This is a more difficult case and we bail.
1669	ArrayRef<int64_t> insertedPos = nextInsertOp.getStaticPosition();
1670	if (isContainedWithin(a: extractPosition, b: insertedPos) ||
1671	intersectsWhereNonNegative(a: extractPosition, b: insertedPos))
1672	return Value();
1673
1674	// Case 5: No intersection, we forward the extract to insertOp.dest().
1675	valueToExtractFrom = nextInsertOp.getDest();
1676	updateStateForNextIteration(v: valueToExtractFrom);
1677	}
1678	// If after all this we can fold, go for it.
1679	return tryToFoldExtractOpInPlace(source: valueToExtractFrom);
1680	}
1681
1682	/// Returns true if the operation has a 0-D vector type operand or result.
1683	static bool hasZeroDimVectors(Operation *op) {
1684	auto hasZeroDimVectorType = [](Type type) -> bool {
1685	auto vecType = dyn_cast<VectorType>(type);
1686	return vecType && vecType.getRank() == 0;
1687	};
1688
1689	return llvm::any_of(Range: op->getOperandTypes(), P: hasZeroDimVectorType) ||
1690	llvm::any_of(Range: op->getResultTypes(), P: hasZeroDimVectorType);
1691	}
1692
1693	/// Fold extractOp with scalar result coming from BroadcastOp or SplatOp.
1694	static Value foldExtractFromBroadcast(ExtractOp extractOp) {
1695	Operation *defOp = extractOp.getVector().getDefiningOp();
1696	if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
1697	return Value();
1698
1699	Value source = defOp->getOperand(idx: 0);
1700	if (extractOp.getType() == source.getType())
1701	return source;
1702	auto getRank = [](Type type) {
1703	return llvm::isa<VectorType>(type) ? llvm::cast<VectorType>(type).getRank()
1704	: 0;
1705	};
1706
1707	// If splat or broadcast from a scalar, just return the source scalar.
1708	unsigned broadcastSrcRank = getRank(source.getType());
1709	if (broadcastSrcRank == 0 && source.getType() == extractOp.getType())
1710	return source;
1711
1712	unsigned extractResultRank = getRank(extractOp.getType());
1713	if (extractResultRank > broadcastSrcRank)
1714	return Value();
1715	// Check that the dimension of the result haven't been broadcasted.
1716	auto extractVecType = llvm::dyn_cast<VectorType>(extractOp.getType());
1717	auto broadcastVecType = llvm::dyn_cast<VectorType>(source.getType());
1718	if (extractVecType && broadcastVecType &&
1719	extractVecType.getShape() !=
1720	broadcastVecType.getShape().take_back(extractResultRank))
1721	return Value();
1722
1723	auto broadcastOp = cast<vector::BroadcastOp>(defOp);
1724	int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank();
1725
1726	// Detect all the positions that come from "dim-1" broadcasting.
1727	// These dimensions correspond to "dim-1" broadcasted dims; set the mathching
1728	// extract position to `0` when extracting from the source operand.
1729	llvm::SetVector<int64_t> broadcastedUnitDims =
1730	broadcastOp.computeBroadcastedUnitDims();
1731	SmallVector<OpFoldResult> extractPos(extractOp.getMixedPosition());
1732	OpBuilder b(extractOp.getContext());
1733	int64_t broadcastRankDiff = broadcastDstRank - broadcastSrcRank;
1734	for (int64_t i = broadcastRankDiff, e = extractPos.size(); i < e; ++i)
1735	if (broadcastedUnitDims.contains(key: i))
1736	extractPos[i] = b.getIndexAttr(0);
1737	// `rankDiff` leading dimensions correspond to new broadcasted dims, drop the
1738	// matching extract position when extracting from the source operand.
1739	int64_t rankDiff = broadcastSrcRank - extractResultRank;
1740	extractPos.erase(CS: extractPos.begin(),
1741	CE: std::next(x: extractPos.begin(), n: extractPos.size() - rankDiff));
1742	// OpBuilder is only used as a helper to build an I64ArrayAttr.
1743	auto [staticPos, dynPos] = decomposeMixedValues(mixedValues: extractPos);
1744	extractOp->setOperands(
1745	llvm::to_vector(llvm::concat<Value>(ValueRange(source), dynPos)));
1746	extractOp.setStaticPosition(staticPos);
1747	return extractOp.getResult();
1748	}
1749
1750	/// Fold extractOp coming from ShuffleOp.
1751	///
1752	/// Example:
1753	///
1754	/// %shuffle = vector.shuffle %a, %b [0, 8, 7, 15]
1755	/// : vector<8xf32>, vector<8xf32>
1756	/// %extract = vector.extract %shuffle[3] : f32 from vector<4xf32>
1757	/// ->
1758	/// %extract = vector.extract %b[7] : f32 from vector<8xf32>
1759	///
1760	static Value foldExtractFromShuffle(ExtractOp extractOp) {
1761	// Dynamic positions are not folded as the resulting code would be more
1762	// complex than the input code.
1763	if (extractOp.hasDynamicPosition())
1764	return Value();
1765
1766	auto shuffleOp = extractOp.getVector().getDefiningOp<ShuffleOp>();
1767	if (!shuffleOp)
1768	return Value();
1769
1770	// TODO: 0-D or multi-dimensional vectors not supported yet.
1771	if (shuffleOp.getResultVectorType().getRank() != 1)
1772	return Value();
1773
1774	int64_t inputVecSize = shuffleOp.getV1().getType().getShape()[0];
1775	auto shuffleMask = shuffleOp.getMask();
1776	int64_t extractIdx = extractOp.getStaticPosition()[0];
1777	int64_t shuffleIdx = shuffleMask[extractIdx];
1778
1779	// Find the shuffled vector to extract from based on the shuffle index.
1780	if (shuffleIdx < inputVecSize) {
1781	extractOp.setOperand(0, shuffleOp.getV1());
1782	extractOp.setStaticPosition({shuffleIdx});
1783	} else {
1784	extractOp.setOperand(0, shuffleOp.getV2());
1785	extractOp.setStaticPosition({shuffleIdx - inputVecSize});
1786	}
1787
1788	return extractOp.getResult();
1789	}
1790
1791	// Fold extractOp with source coming from ShapeCast op.
1792	static Value foldExtractFromShapeCast(ExtractOp extractOp) {
1793	// TODO: Canonicalization for dynamic position not implemented yet.
1794	if (extractOp.hasDynamicPosition())
1795	return Value();
1796
1797	auto shapeCastOp = extractOp.getVector().getDefiningOp<vector::ShapeCastOp>();
1798	if (!shapeCastOp)
1799	return Value();
1800
1801	// Get the nth dimension size starting from lowest dimension.
1802	auto getDimReverse = [](VectorType type, int64_t n) {
1803	return type.getShape().take_back(n + 1).front();
1804	};
1805	int64_t destinationRank =
1806	llvm::isa<VectorType>(extractOp.getType())
1807	? llvm::cast<VectorType>(extractOp.getType()).getRank()
1808	: 0;
1809	if (destinationRank > shapeCastOp.getSourceVectorType().getRank())
1810	return Value();
1811	if (destinationRank > 0) {
1812	auto destinationType =
1813	llvm::cast<VectorType>(extractOp.getResult().getType());
1814	for (int64_t i = 0; i < destinationRank; i++) {
1815	// The lowest dimension of the destination must match the lowest
1816	// dimension of the shapecast op source.
1817	// TODO: This case could be support in a canonicalization pattern.
1818	if (getDimReverse(shapeCastOp.getSourceVectorType(), i) !=
1819	getDimReverse(destinationType, i))
1820	return Value();
1821	}
1822	}
1823	// Extract the strides associated with the extract op vector source. Then use
1824	// this to calculate a linearized position for the extract.
1825	SmallVector<int64_t> extractedPos(extractOp.getStaticPosition());
1826	std::reverse(first: extractedPos.begin(), last: extractedPos.end());
1827	SmallVector<int64_t, 4> strides;
1828	int64_t stride = 1;
1829	for (int64_t i = 0, e = extractedPos.size(); i < e; i++) {
1830	strides.push_back(Elt: stride);
1831	stride *=
1832	getDimReverse(extractOp.getSourceVectorType(), i + destinationRank);
1833	}
1834
1835	int64_t position = linearize(offsets: extractedPos, basis: strides);
1836	// Then extract the strides associated to the shapeCast op vector source and
1837	// delinearize the position using those strides.
1838	SmallVector<int64_t, 4> newStrides;
1839	int64_t numDimension =
1840	shapeCastOp.getSourceVectorType().getRank() - destinationRank;
1841	stride = 1;
1842	for (int64_t i = 0; i < numDimension; i++) {
1843	newStrides.push_back(Elt: stride);
1844	stride *=
1845	getDimReverse(shapeCastOp.getSourceVectorType(), i + destinationRank);
1846	}
1847	std::reverse(first: newStrides.begin(), last: newStrides.end());
1848	SmallVector<int64_t, 4> newPosition = delinearize(linearIndex: position, strides: newStrides);
1849	// OpBuilder is only used as a helper to build an I64ArrayAttr.
1850	OpBuilder b(extractOp.getContext());
1851	extractOp.setStaticPosition(newPosition);
1852	extractOp.setOperand(0, shapeCastOp.getSource());
1853	return extractOp.getResult();
1854	}
1855
1856	/// Fold an ExtractOp from ExtractStridedSliceOp.
1857	static Value foldExtractFromExtractStrided(ExtractOp extractOp) {
1858	// TODO: Canonicalization for dynamic position not implemented yet.
1859	if (extractOp.hasDynamicPosition())
1860	return Value();
1861
1862	auto extractStridedSliceOp =
1863	extractOp.getVector().getDefiningOp<vector::ExtractStridedSliceOp>();
1864	if (!extractStridedSliceOp)
1865	return Value();
1866
1867	// 0-D vectors not supported.
1868	assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
1869	if (hasZeroDimVectors(extractStridedSliceOp))
1870	return Value();
1871
1872	// Return if 'extractStridedSliceOp' has non-unit strides.
1873	if (extractStridedSliceOp.hasNonUnitStrides())
1874	return Value();
1875
1876	// Trim offsets for dimensions fully extracted.
1877	auto sliceOffsets =
1878	extractVector<int64_t>(extractStridedSliceOp.getOffsets());
1879	while (!sliceOffsets.empty()) {
1880	size_t lastOffset = sliceOffsets.size() - 1;
1881	if (sliceOffsets.back() != 0 ||
1882	extractStridedSliceOp.getType().getDimSize(lastOffset) !=
1883	extractStridedSliceOp.getSourceVectorType().getDimSize(lastOffset))
1884	break;
1885	sliceOffsets.pop_back();
1886	}
1887	unsigned destinationRank = 0;
1888	if (auto vecType = llvm::dyn_cast<VectorType>(extractOp.getType()))
1889	destinationRank = vecType.getRank();
1890	// The dimensions of the result need to be untouched by the
1891	// extractStridedSlice op.
1892	if (destinationRank > extractStridedSliceOp.getSourceVectorType().getRank() -
1893	sliceOffsets.size())
1894	return Value();
1895
1896	SmallVector<int64_t> extractedPos(extractOp.getStaticPosition());
1897	assert(extractedPos.size() >= sliceOffsets.size());
1898	for (size_t i = 0, e = sliceOffsets.size(); i < e; i++)
1899	extractedPos[i] = extractedPos[i] + sliceOffsets[i];
1900	extractOp.getVectorMutable().assign(extractStridedSliceOp.getVector());
1901
1902	// OpBuilder is only used as a helper to build an I64ArrayAttr.
1903	OpBuilder b(extractOp.getContext());
1904	extractOp.setStaticPosition(extractedPos);
1905	return extractOp.getResult();
1906	}
1907
1908	/// Fold extract_op fed from a chain of insertStridedSlice ops.
1909	static Value foldExtractStridedOpFromInsertChain(ExtractOp extractOp) {
1910	// TODO: Canonicalization for dynamic position not implemented yet.
1911	if (extractOp.hasDynamicPosition())
1912	return Value();
1913
1914	int64_t destinationRank =
1915	llvm::isa<VectorType>(extractOp.getType())
1916	? llvm::cast<VectorType>(extractOp.getType()).getRank()
1917	: 0;
1918	auto insertOp = extractOp.getVector().getDefiningOp<InsertStridedSliceOp>();
1919	if (!insertOp)
1920	return Value();
1921
1922	// 0-D vectors not supported.
1923	assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
1924	if (hasZeroDimVectors(insertOp))
1925	return Value();
1926
1927	while (insertOp) {
1928	int64_t insertRankDiff = insertOp.getDestVectorType().getRank() -
1929	insertOp.getSourceVectorType().getRank();
1930	if (destinationRank > insertOp.getSourceVectorType().getRank())
1931	return Value();
1932	auto insertOffsets = extractVector<int64_t>(insertOp.getOffsets());
1933	ArrayRef<int64_t> extractOffsets = extractOp.getStaticPosition();
1934
1935	if (llvm::any_of(insertOp.getStrides(), [](Attribute attr) {
1936	return llvm::cast<IntegerAttr>(attr).getInt() != 1;
1937	}))
1938	return Value();
1939	bool disjoint = false;
1940	SmallVector<int64_t, 4> offsetDiffs;
1941	for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
1942	int64_t start = insertOffsets[dim];
1943	int64_t size =
1944	(dim < insertRankDiff)
1945	? 1
1946	: insertOp.getSourceVectorType().getDimSize(dim - insertRankDiff);
1947	int64_t end = start + size;
1948	int64_t offset = extractOffsets[dim];
1949	// Check if the start of the extract offset is in the interval inserted.
1950	if (start <= offset && offset < end) {
1951	if (dim >= insertRankDiff)
1952	offsetDiffs.push_back(Elt: offset - start);
1953	continue;
1954	}
1955	disjoint = true;
1956	break;
1957	}
1958	// The extract element chunk overlap with the vector inserted.
1959	if (!disjoint) {
1960	// If any of the inner dimensions are only partially inserted we have a
1961	// partial overlap.
1962	int64_t srcRankDiff =
1963	insertOp.getSourceVectorType().getRank() - destinationRank;
1964	for (int64_t i = 0; i < destinationRank; i++) {
1965	if (insertOp.getSourceVectorType().getDimSize(i + srcRankDiff) !=
1966	insertOp.getDestVectorType().getDimSize(i + srcRankDiff +
1967	insertRankDiff))
1968	return Value();
1969	}
1970	extractOp.getVectorMutable().assign(insertOp.getValueToStore());
1971	// OpBuilder is only used as a helper to build an I64ArrayAttr.
1972	OpBuilder b(extractOp.getContext());
1973	extractOp.setStaticPosition(offsetDiffs);
1974	return extractOp.getResult();
1975	}
1976	// If the chunk extracted is disjoint from the chunk inserted, keep
1977	// looking in the insert chain.
1978	insertOp = insertOp.getDest().getDefiningOp<InsertStridedSliceOp>();
1979	}
1980	return Value();
1981	}
1982
1983	/// Try to fold the extraction of a scalar from a vector defined by
1984	/// vector.from_elements. E.g.:
1985	///
1986	/// %0 = vector.from_elements %a, %b : vector<2xf32>
1987	/// %1 = vector.extract %0[0] : f32 from vector<2xf32>
1988	/// ==> fold to %a
1989	static Value foldScalarExtractFromFromElements(ExtractOp extractOp) {
1990	// Dynamic extractions cannot be folded.
1991	if (extractOp.hasDynamicPosition())
1992	return {};
1993
1994	// Look for extract(from_elements).
1995	auto fromElementsOp = extractOp.getVector().getDefiningOp<FromElementsOp>();
1996	if (!fromElementsOp)
1997	return {};
1998
1999	// Scalable vectors are not supported.
2000	auto vecType = llvm::cast<VectorType>(fromElementsOp.getType());
2001	if (vecType.isScalable())
2002	return {};
2003
2004	// Only extractions of scalars are supported.
2005	int64_t rank = vecType.getRank();
2006	ArrayRef<int64_t> indices = extractOp.getStaticPosition();
2007	if (extractOp.getType() != vecType.getElementType())
2008	return {};
2009	assert(static_cast<int64_t>(indices.size()) == rank &&
2010	"unexpected number of indices");
2011
2012	// Compute flattened/linearized index and fold to operand.
2013	int flatIndex = 0;
2014	int stride = 1;
2015	for (int i = rank - 1; i >= 0; --i) {
2016	flatIndex += indices[i] * stride;
2017	stride *= vecType.getDimSize(i);
2018	}
2019	return fromElementsOp.getElements()[flatIndex];
2020	}
2021
2022	/// If the dynamic indices of `extractOp` or `insertOp` are in fact constants,
2023	/// then fold it.
2024	template <typename OpType, typename AdaptorType>
2025	static Value extractInsertFoldConstantOp(OpType op, AdaptorType adaptor,
2026	SmallVectorImpl<Value> &operands) {
2027	std::vector<int64_t> staticPosition = op.getStaticPosition().vec();
2028	OperandRange dynamicPosition = op.getDynamicPosition();
2029	ArrayRef<Attribute> dynamicPositionAttr = adaptor.getDynamicPosition();
2030	ArrayRef<int64_t> vectorShape;
2031	if constexpr (std::is_same_v<OpType, ExtractOp>)
2032	vectorShape = op.getSourceVectorType().getShape();
2033	else
2034	vectorShape = op.getDestVectorType().getShape();
2035
2036	// If the dynamic operands is empty, it is returned directly.
2037	if (!dynamicPosition.size())
2038	return {};
2039
2040	// `index` is used to iterate over the `dynamicPosition`.
2041	unsigned index = 0;
2042
2043	// `opChange` is a flag. If it is true, it means to update `op` in place.
2044	bool opChange = false;
2045	for (unsigned i = 0, e = staticPosition.size(); i < e; ++i) {
2046	if (!ShapedType::isDynamic(staticPosition[i]))
2047	continue;
2048	Attribute positionAttr = dynamicPositionAttr[index];
2049	Value position = dynamicPosition[index++];
2050	if (auto attr = mlir::dyn_cast_if_present<IntegerAttr>(positionAttr)) {
2051	int64_t value = attr.getInt();
2052	// Do not fold if the value is out of bounds (-1 signifies a poison
2053	// value rather than OOB index).
2054	if (value >= -1 && value < vectorShape[i]) {
2055	staticPosition[i] = attr.getInt();
2056	opChange = true;
2057	continue;
2058	}
2059	}
2060	operands.push_back(Elt: position);
2061	}
2062
2063	if (opChange) {
2064	op.setStaticPosition(staticPosition);
2065	op.getOperation()->setOperands(operands);
2066	return op.getResult();
2067	}
2068	return {};
2069	}
2070
2071	/// Fold an insert or extract operation into an poison value when a poison index
2072	/// is found at any dimension of the static position.
2073	static Attribute foldPoisonIndexInsertExtractOp(MLIRContext *context,
2074	ArrayRef<int64_t> staticPos,
2075	int64_t poisonVal) {
2076	if (!is_contained(Range&: staticPos, Element: poisonVal))
2077	return {};
2078
2079	return ub::PoisonAttr::get(context);
2080	}
2081
2082	/// Fold a vector extract from is a poison source.
2083	static Attribute foldPoisonSrcExtractOp(Attribute srcAttr) {
2084	if (isa_and_nonnull<ub::PoisonAttr>(srcAttr))
2085	return srcAttr;
2086
2087	return {};
2088	}
2089
2090	/// Fold a vector extract extracting from a DenseElementsAttr.
2091	static Attribute foldDenseElementsAttrSrcExtractOp(ExtractOp extractOp,
2092	Attribute srcAttr) {
2093	auto denseAttr = dyn_cast_if_present<DenseElementsAttr>(Val&: srcAttr);
2094	if (!denseAttr) {
2095	return {};
2096	}
2097
2098	if (denseAttr.isSplat()) {
2099	Attribute newAttr = denseAttr.getSplatValue<Attribute>();
2100	if (auto vecDstType = dyn_cast<VectorType>(extractOp.getType()))
2101	newAttr = DenseElementsAttr::get(vecDstType, newAttr);
2102	return newAttr;
2103	}
2104
2105	auto vecTy = cast<VectorType>(extractOp.getSourceVectorType());
2106	if (vecTy.isScalable())
2107	return {};
2108
2109	if (extractOp.hasDynamicPosition()) {
2110	return {};
2111	}
2112
2113	// Materializing subsets of a large constant array can generally lead to
2114	// explosion in IR size because of different combination of subsets that
2115	// can exist. However, vector.extract is a restricted form of subset
2116	// extract where you can only extract non-overlapping (or the same) subset for
2117	// a given rank of the subset. Because of this property, the IR size can only
2118	// increase at most by `rank * size(array)` from a single constant array being
2119	// extracted by multiple extracts.
2120
2121	// Calculate the linearized position of the continuous chunk of elements to
2122	// extract.
2123	SmallVector<int64_t> completePositions(vecTy.getRank(), 0);
2124	copy(extractOp.getStaticPosition(), completePositions.begin());
2125	int64_t startPos =
2126	linearize(completePositions, computeStrides(vecTy.getShape()));
2127	auto denseValuesBegin = denseAttr.value_begin<TypedAttr>() + startPos;
2128
2129	TypedAttr newAttr;
2130	if (auto resVecTy = dyn_cast<VectorType>(extractOp.getType())) {
2131	SmallVector<Attribute> elementValues(
2132	denseValuesBegin, denseValuesBegin + resVecTy.getNumElements());
2133	newAttr = DenseElementsAttr::get(resVecTy, elementValues);
2134	} else {
2135	newAttr = *denseValuesBegin;
2136	}
2137
2138	return newAttr;
2139	}
2140
2141	OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
2142	// Fold "vector.extract %v[] : vector<2x2xf32> from vector<2x2xf32>" to %v.
2143	// Note: Do not fold "vector.extract %v[] : f32 from vector<f32>" (type
2144	// mismatch).
2145	if (getNumIndices() == 0 && getVector().getType() == getResult().getType())
2146	return getVector();
2147	if (auto res = foldPoisonSrcExtractOp(adaptor.getVector()))
2148	return res;
2149	// Fold `arith.constant` indices into the `vector.extract` operation. Make
2150	// sure that patterns requiring constant indices are added after this fold.
2151	SmallVector<Value> operands = {getVector()};
2152	if (auto val = extractInsertFoldConstantOp(*this, adaptor, operands))
2153	return val;
2154	if (auto res = foldPoisonIndexInsertExtractOp(
2155	getContext(), adaptor.getStaticPosition(), kPoisonIndex))
2156	return res;
2157	if (auto res = foldDenseElementsAttrSrcExtractOp(*this, adaptor.getVector()))
2158	return res;
2159	if (succeeded(foldExtractOpFromExtractChain(*this)))
2160	return getResult();
2161	if (auto res = ExtractFromInsertTransposeChainState(*this).fold())
2162	return res;
2163	if (auto res = foldExtractFromBroadcast(*this))
2164	return res;
2165	if (auto res = foldExtractFromShuffle(*this))
2166	return res;
2167	if (auto res = foldExtractFromShapeCast(*this))
2168	return res;
2169	if (auto val = foldExtractFromExtractStrided(*this))
2170	return val;
2171	if (auto val = foldExtractStridedOpFromInsertChain(*this))
2172	return val;
2173	if (auto val = foldScalarExtractFromFromElements(*this))
2174	return val;
2175	return OpFoldResult();
2176	}
2177
2178	namespace {
2179
2180	// Pattern to rewrite a ExtractOp(Broadcast) -> Broadcast.
2181	class ExtractOpFromBroadcast final : public OpRewritePattern<ExtractOp> {
2182	public:
2183	using OpRewritePattern::OpRewritePattern;
2184
2185	LogicalResult matchAndRewrite(ExtractOp extractOp,
2186	PatternRewriter &rewriter) const override {
2187	Operation *defOp = extractOp.getVector().getDefiningOp();
2188	if (!defOp || !isa<vector::BroadcastOp, SplatOp>(defOp))
2189	return failure();
2190
2191	Value source = defOp->getOperand(idx: 0);
2192	if (extractOp.getType() == source.getType())
2193	return failure();
2194	auto getRank = [](Type type) {
2195	return llvm::isa<VectorType>(type)
2196	? llvm::cast<VectorType>(type).getRank()
2197	: 0;
2198	};
2199	unsigned broadcastSrcRank = getRank(source.getType());
2200	unsigned extractResultRank = getRank(extractOp.getType());
2201	// We only consider the case where the rank of the source is less than or
2202	// equal to the rank of the extract dst. The other cases are handled in the
2203	// folding patterns.
2204	if (extractResultRank < broadcastSrcRank)
2205	return failure();
2206	// For scalar result, the input can only be a rank-0 vector, which will
2207	// be handled by the folder.
2208	if (extractResultRank == 0)
2209	return failure();
2210
2211	rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
2212	extractOp, extractOp.getType(), source);
2213	return success();
2214	}
2215	};
2216
2217	// Pattern to rewrite a ExtractOp(CreateMask) -> CreateMask.
2218	class ExtractOpFromCreateMask final : public OpRewritePattern<ExtractOp> {
2219	public:
2220	using OpRewritePattern::OpRewritePattern;
2221
2222	LogicalResult matchAndRewrite(ExtractOp extractOp,
2223	PatternRewriter &rewriter) const override {
2224	auto createMaskOp =
2225	extractOp.getVector().getDefiningOp<vector::CreateMaskOp>();
2226	if (!createMaskOp)
2227	return failure();
2228
2229	VectorType extractedMaskType =
2230	llvm::dyn_cast<VectorType>(extractOp.getResult().getType());
2231
2232	if (!extractedMaskType)
2233	return failure();
2234
2235	auto maskOperands = createMaskOp.getOperands();
2236	ArrayRef<int64_t> extractOpPos = extractOp.getStaticPosition();
2237	VectorType maskType = createMaskOp.getVectorType();
2238
2239	bool containsUnknownDims = false;
2240	bool allFalse = getMaskFormat(createMaskOp) == MaskFormat::AllFalse;
2241
2242	for (size_t dimIdx = 0; !allFalse && dimIdx < extractOpPos.size();
2243	dimIdx++) {
2244	int64_t pos = extractOpPos[dimIdx];
2245	Value operand = maskOperands[dimIdx];
2246	auto constantOp = operand.getDefiningOp<arith::ConstantOp>();
2247	if (!constantOp) {
2248	// Bounds of this dim unknown.
2249	containsUnknownDims = true;
2250	continue;
2251	}
2252
2253	int64_t createMaskBound =
2254	llvm::cast<IntegerAttr>(constantOp.getValue()).getInt();
2255
2256	if (pos != ShapedType::kDynamic) {
2257	// If any position is outside the range from the `create_mask`, then the
2258	// extracted mask will be all-false.
2259	allFalse |= pos >= createMaskBound;
2260	} else if (createMaskBound < maskType.getDimSize(dimIdx)) {
2261	// This dim is not all-true and since this is a dynamic index we don't
2262	// know if the extraction is within the true or false region.
2263	// Note: Zero dims have already handled via getMaskFormat().
2264	containsUnknownDims = true;
2265	}
2266	}
2267
2268	if (allFalse) {
2269	rewriter.replaceOpWithNewOp<arith::ConstantOp>(
2270	extractOp, DenseElementsAttr::get(extractedMaskType, false));
2271	} else if (!containsUnknownDims) {
2272	rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
2273	extractOp, extractedMaskType,
2274	maskOperands.drop_front(extractOpPos.size()));
2275	} else {
2276	return failure();
2277	}
2278	return success();
2279	}
2280	};
2281
2282	// Folds extract(shape_cast(..)) into shape_cast when the total element count
2283	// does not change.
2284	LogicalResult foldExtractFromShapeCastToShapeCast(ExtractOp extractOp,
2285	PatternRewriter &rewriter) {
2286	auto castOp = extractOp.getVector().getDefiningOp<ShapeCastOp>();
2287	if (!castOp)
2288	return failure();
2289
2290	VectorType sourceType = castOp.getSourceVectorType();
2291	auto targetType = dyn_cast<VectorType>(extractOp.getResult().getType());
2292	if (!targetType)
2293	return failure();
2294
2295	if (sourceType.getNumElements() != targetType.getNumElements())
2296	return failure();
2297
2298	rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(extractOp, targetType,
2299	castOp.getSource());
2300	return success();
2301	}
2302
2303	/// Try to canonicalize the extraction of a subvector from a vector defined by
2304	/// vector.from_elements. E.g.:
2305	///
2306	/// %0 = vector.from_elements %a, %b, %a, %a : vector<2x2xf32>
2307	/// %1 = vector.extract %0[0] : vector<2xf32> from vector<2x2xf32>
2308	/// ==> canonicalize to vector.from_elements %a, %b : vector<2xf32>
2309	LogicalResult foldExtractFromFromElements(ExtractOp extractOp,
2310	PatternRewriter &rewriter) {
2311	// Dynamic positions are not supported.
2312	if (extractOp.hasDynamicPosition())
2313	return failure();
2314
2315	// Scalar extracts are handled by the folder.
2316	auto resultType = dyn_cast<VectorType>(extractOp.getType());
2317	if (!resultType)
2318	return failure();
2319
2320	// Look for extracts from a from_elements op.
2321	auto fromElementsOp = extractOp.getVector().getDefiningOp<FromElementsOp>();
2322	if (!fromElementsOp)
2323	return failure();
2324	VectorType inputType = fromElementsOp.getType();
2325
2326	// Scalable vectors are not supported.
2327	if (resultType.isScalable() || inputType.isScalable())
2328	return failure();
2329
2330	// Compute the position of first extracted element and flatten/linearize the
2331	// position.
2332	SmallVector<int64_t> firstElementPos =
2333	llvm::to_vector(extractOp.getStaticPosition());
2334	firstElementPos.append(/*NumInputs=*/resultType.getRank(), /*Elt=*/0);
2335	int flatIndex = 0;
2336	int stride = 1;
2337	for (int64_t i = inputType.getRank() - 1; i >= 0; --i) {
2338	flatIndex += firstElementPos[i] * stride;
2339	stride *= inputType.getDimSize(i);
2340	}
2341
2342	// Replace the op with a smaller from_elements op.
2343	rewriter.replaceOpWithNewOp<FromElementsOp>(
2344	extractOp, resultType,
2345	fromElementsOp.getElements().slice(flatIndex,
2346	resultType.getNumElements()));
2347	return success();
2348	}
2349
2350	} // namespace
2351
2352	void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
2353	MLIRContext *context) {
2354	results.add<ExtractOpFromBroadcast, ExtractOpFromCreateMask>(context);
2355	results.add(foldExtractFromShapeCastToShapeCast);
2356	results.add(foldExtractFromFromElements);
2357	}
2358
2359	static void populateFromInt64AttrArray(ArrayAttr arrayAttr,
2360	SmallVectorImpl<int64_t> &results) {
2361	for (auto attr : arrayAttr)
2362	results.push_back(llvm::cast<IntegerAttr>(attr).getInt());
2363	}
2364
2365	//===----------------------------------------------------------------------===//
2366	// FmaOp
2367	//===----------------------------------------------------------------------===//
2368
2369	std::optional<SmallVector<int64_t, 4>> FMAOp::getShapeForUnroll() {
2370	return llvm::to_vector<4>(getVectorType().getShape());
2371	}
2372
2373	//===----------------------------------------------------------------------===//
2374	// FromElementsOp
2375	//===----------------------------------------------------------------------===//
2376
2377	/// Rewrite a vector.from_elements into a vector.splat if all elements are the
2378	/// same SSA value. E.g.:
2379	///
2380	/// %0 = vector.from_elements %a, %a, %a : vector<3xf32>
2381	/// ==> rewrite to vector.splat %a : vector<3xf32>
2382	static LogicalResult rewriteFromElementsAsSplat(FromElementsOp fromElementsOp,
2383	PatternRewriter &rewriter) {
2384	if (!llvm::all_equal(fromElementsOp.getElements()))
2385	return failure();
2386	rewriter.replaceOpWithNewOp<SplatOp>(fromElementsOp, fromElementsOp.getType(),
2387	fromElementsOp.getElements().front());
2388	return success();
2389	}
2390
2391	/// Rewrite from_elements on multiple scalar extracts as a shape_cast
2392	/// on a single extract. Example:
2393	/// %0 = vector.extract %source[0, 0] : i8 from vector<2x2xi8>
2394	/// %1 = vector.extract %source[0, 1] : i8 from vector<2x2xi8>
2395	/// %2 = vector.from_elements %0, %1 : vector<2xi8>
2396	///
2397	/// becomes
2398	/// %1 = vector.extract %source[0] : vector<1x2xi8> from vector<2x2xi8>
2399	/// %2 = vector.shape_cast %1 : vector<1x2xi8> to vector<2xi8>
2400	///
2401	/// The requirements for this to be valid are
2402	///
2403	/// i) The elements are extracted from the same vector (%source).
2404	///
2405	/// ii) The elements form a suffix of %source. Specifically, the number
2406	/// of elements is the same as the product of the last N dimension sizes
2407	/// of %source, for some N.
2408	///
2409	/// iii) The elements are extracted contiguously in ascending order.
2410
2411	class FromElementsToShapeCast : public OpRewritePattern<FromElementsOp> {
2412
2413	using OpRewritePattern::OpRewritePattern;
2414
2415	LogicalResult matchAndRewrite(FromElementsOp fromElements,
2416	PatternRewriter &rewriter) const override {
2417
2418	// Handled by `rewriteFromElementsAsSplat`
2419	if (fromElements.getType().getNumElements() == 1)
2420	return failure();
2421
2422	// The common source that all elements are extracted from, if one exists.
2423	TypedValue<VectorType> source;
2424	// The position of the combined extract operation, if one is created.
2425	ArrayRef<int64_t> combinedPosition;
2426	// The expected index of extraction of the current element in the loop, if
2427	// elements are extracted contiguously in ascending order.
2428	SmallVector<int64_t> expectedPosition;
2429
2430	for (auto [insertIndex, element] :
2431	llvm::enumerate(fromElements.getElements())) {
2432
2433	// Check that the element is from a vector.extract operation.
2434	auto extractOp =
2435	dyn_cast_if_present<vector::ExtractOp>(element.getDefiningOp());
2436	if (!extractOp) {
2437	return rewriter.notifyMatchFailure(fromElements,
2438	"element not from vector.extract");
2439	}
2440
2441	// Check condition (i) by checking that all elements have the same source
2442	// as the first element.
2443	if (insertIndex == 0) {
2444	source = extractOp.getVector();
2445	} else if (extractOp.getVector() != source) {
2446	return rewriter.notifyMatchFailure(fromElements,
2447	"element from different vector");
2448	}
2449
2450	ArrayRef<int64_t> position = extractOp.getStaticPosition();
2451	int64_t rank = position.size();
2452	assert(rank == source.getType().getRank() &&
2453	"scalar extract must have full rank position");
2454
2455	// Check condition (ii) by checking that the position that the first
2456	// element is extracted from has sufficient trailing 0s. For example, in
2457	//
2458	// %elm0 = vector.extract %source[1, 0, 0] : i8 from vector<2x3x4xi8>
2459	// [...]
2460	// %elms = vector.from_elements %elm0, [...] : vector<12xi8>
2461	//
2462	// The 2 trailing 0s in the position of extraction of %elm0 cover 3*4 = 12
2463	// elements, which is the number of elements of %n, so this is valid.
2464	if (insertIndex == 0) {
2465	const int64_t numElms = fromElements.getType().getNumElements();
2466	int64_t numSuffixElms = 1;
2467	int64_t index = rank;
2468	while (index > 0 && position[index - 1] == 0 &&
2469	numSuffixElms < numElms) {
2470	numSuffixElms *= source.getType().getDimSize(index - 1);
2471	--index;
2472	}
2473	if (numSuffixElms != numElms) {
2474	return rewriter.notifyMatchFailure(
2475	fromElements, "elements do not form a suffix of source");
2476	}
2477	expectedPosition = llvm::to_vector(position);
2478	combinedPosition = position.drop_back(rank - index);
2479	}
2480
2481	// Check condition (iii).
2482	else if (expectedPosition != position) {
2483	return rewriter.notifyMatchFailure(
2484	fromElements, "elements not in ascending order (static order)");
2485	}
2486	increment(expectedPosition, source.getType().getShape());
2487	}
2488
2489	auto extracted = rewriter.createOrFold<vector::ExtractOp>(
2490	fromElements.getLoc(), source, combinedPosition);
2491
2492	rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
2493	fromElements, fromElements.getType(), extracted);
2494
2495	return success();
2496	}
2497
2498	/// Increments n-D `indices` by 1 starting from the innermost dimension.
2499	static void increment(MutableArrayRef<int64_t> indices,
2500	ArrayRef<int64_t> shape) {
2501	for (int dim : llvm::reverse(C: llvm::seq<int>(Begin: 0, End: indices.size()))) {
2502	indices[dim] += 1;
2503	if (indices[dim] < shape[dim])
2504	break;
2505	indices[dim] = 0;
2506	}
2507	}
2508	};
2509
2510	void FromElementsOp::getCanonicalizationPatterns(RewritePatternSet &results,
2511	MLIRContext *context) {
2512	results.add(rewriteFromElementsAsSplat);
2513	results.add<FromElementsToShapeCast>(context);
2514	}
2515
2516	//===----------------------------------------------------------------------===//
2517	// BroadcastOp
2518	//===----------------------------------------------------------------------===//
2519
2520	void BroadcastOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
2521	SetIntRangeFn setResultRanges) {
2522	setResultRanges(getResult(), argRanges.front());
2523	}
2524
2525	std::optional<SmallVector<int64_t, 4>> BroadcastOp::getShapeForUnroll() {
2526	return llvm::to_vector<4>(getResultVectorType().getShape());
2527	}
2528
2529	/// Return the dimensions of the result vector that were formerly ones in the
2530	/// source tensor and thus correspond to "dim-1" broadcasting.
2531	static llvm::SetVector<int64_t>
2532	computeBroadcastedUnitDims(ArrayRef<int64_t> srcShape,
2533	ArrayRef<int64_t> dstShape) {
2534	int64_t rankDiff = dstShape.size() - srcShape.size();
2535	int64_t dstDim = rankDiff;
2536	llvm::SetVector<int64_t> res;
2537	for (auto [s1, s2] :
2538	llvm::zip_equal(t&: srcShape, u: dstShape.drop_front(N: rankDiff))) {
2539	if (s1 != s2) {
2540	assert(s1 == 1 && "expected \"dim-1\" broadcasting");
2541	res.insert(X: dstDim);
2542	}
2543	++dstDim;
2544	}
2545	return res;
2546	}
2547
2548	llvm::SetVector<int64_t> BroadcastOp::computeBroadcastedUnitDims() {
2549	// Scalar broadcast is without any unit dim broadcast.
2550	auto srcVectorType = llvm::dyn_cast<VectorType>(getSourceType());
2551	if (!srcVectorType)
2552	return {};
2553	return ::computeBroadcastedUnitDims(srcVectorType.getShape(),
2554	getResultVectorType().getShape());
2555	}
2556
2557	/// Broadcast `value` to a vector of `dstShape`, knowing that exactly the
2558	/// `broadcastedDims` dimensions in the dstShape are broadcasted.
2559	/// This requires (and asserts) that the broadcast is free of "dim-1"
2560	/// broadcasting.
2561	/// Since vector.broadcast only allows expanding leading dimensions, an extra
2562	/// vector.transpose may be inserted to make the broadcast possible.
2563	/// `value`, `dstShape` and `broadcastedDims` must be properly specified or
2564	/// the helper will assert. This means:
2565	/// 1. `dstShape` must not be empty.
2566	/// 2. `broadcastedDims` must be confined to [0 .. rank(value.getVectorType)]
2567	/// 2. `dstShape` trimmed of the dimensions specified in `broadcastedDims`
2568	// must match the `value` shape.
2569	Value BroadcastOp::createOrFoldBroadcastOp(
2570	OpBuilder &b, Value value, ArrayRef<int64_t> dstShape,
2571	const llvm::SetVector<int64_t> &broadcastedDims) {
2572	assert(!dstShape.empty() && "unexpected empty dst shape");
2573
2574	// Well-formedness check.
2575	SmallVector<int64_t> checkShape;
2576	for (int i = 0, e = dstShape.size(); i < e; ++i) {
2577	if (broadcastedDims.contains(i))
2578	continue;
2579	checkShape.push_back(dstShape[i]);
2580	}
2581	assert(broadcastedDims.size() == dstShape.size() - checkShape.size() &&
2582	"ill-formed broadcastedDims contains values not confined to "
2583	"destVectorShape");
2584
2585	Location loc = value.getLoc();
2586	Type elementType = getElementTypeOrSelf(value.getType());
2587	VectorType srcVectorType = llvm::dyn_cast<VectorType>(value.getType());
2588	VectorType dstVectorType = VectorType::get(dstShape, elementType);
2589
2590	// Step 2. If scalar -> dstShape broadcast, just do it.
2591	if (!srcVectorType) {
2592	assert(checkShape.empty() &&
2593	"ill-formed createOrFoldBroadcastOp arguments");
2594	return b.createOrFold<vector::BroadcastOp>(loc, dstVectorType, value);
2595	}
2596
2597	assert(srcVectorType.getShape().equals(checkShape) &&
2598	"ill-formed createOrFoldBroadcastOp arguments");
2599
2600	// Step 3. Since vector.broadcast only allows creating leading dims,
2601	// vector -> dstShape broadcast may require a transpose.
2602	// Traverse the dims in order and construct:
2603	// 1. The leading entries of the broadcastShape that is guaranteed to be
2604	// achievable by a simple broadcast.
2605	// 2. The induced permutation for the subsequent vector.transpose that will
2606	// bring us from `broadcastShape` back to he desired `dstShape`.
2607	// If the induced permutation is not the identity, create a vector.transpose.
2608	SmallVector<int64_t> broadcastShape, permutation(dstShape.size(), -1);
2609	broadcastShape.reserve(dstShape.size());
2610	// Consider the example:
2611	// srcShape = 2x4
2612	// dstShape = 1x2x3x4x5
2613	// broadcastedDims = [0, 2, 4]
2614	//
2615	// We want to build:
2616	// broadcastShape = 1x3x5x2x4
2617	// permutation = [0, 2, 4, 1, 3]
2618	// ---V--- -----V-----
2619	// leading broadcast part src shape part
2620	//
2621	// Note that the trailing dims of broadcastShape are exactly the srcShape
2622	// by construction.
2623	// nextSrcShapeDim is used to keep track of where in the permutation the
2624	// "src shape part" occurs.
2625	int64_t nextSrcShapeDim = broadcastedDims.size();
2626	for (int64_t i = 0, e = dstShape.size(); i < e; ++i) {
2627	if (broadcastedDims.contains(i)) {
2628	// 3.a. For each dim in the dst shape, if it is a broadcasted dim,
2629	// bring it to the head of the broadcastShape.
2630	// It will need to be permuted back from `broadcastShape.size() - 1` into
2631	// position `i`.
2632	broadcastShape.push_back(dstShape[i]);
2633	permutation[i] = broadcastShape.size() - 1;
2634	} else {
2635	// 3.b. Otherwise, the dim is not broadcasted, it comes from the src
2636	// shape and needs to be permuted into position `i`.
2637	// Don't touch `broadcastShape` here, the whole srcShape will be
2638	// appended after.
2639	permutation[i] = nextSrcShapeDim++;
2640	}
2641	}
2642	// 3.c. Append the srcShape.
2643	llvm::append_range(broadcastShape, srcVectorType.getShape());
2644
2645	// Ensure there are no "dim-1" broadcasts.
2646	assert(::computeBroadcastedUnitDims(srcVectorType.getShape(), broadcastShape)
2647	.empty() &&
2648	"unexpected \"dim-1\" broadcast");
2649
2650	VectorType broadcastType = VectorType::get(broadcastShape, elementType);
2651	assert(vector::isBroadcastableTo(value.getType(), broadcastType) ==
2652	vector::BroadcastableToResult::Success &&
2653	"must be broadcastable");
2654	Value res = b.createOrFold<vector::BroadcastOp>(loc, broadcastType, value);
2655	// Step 4. If we find any dimension that indeed needs to be permuted,
2656	// immediately return a new vector.transpose.
2657	for (int64_t i = 0, e = permutation.size(); i < e; ++i)
2658	if (permutation[i] != i)
2659	return b.createOrFold<vector::TransposeOp>(loc, res, permutation);
2660	// Otherwise return res.
2661	return res;
2662	}
2663
2664	BroadcastableToResult mlir::vector::isBroadcastableTo(
2665	Type srcType, VectorType dstVectorType,
2666	std::pair<VectorDim, VectorDim> *mismatchingDims) {
2667	// Broadcast scalar to vector of the same element type.
2668	if (srcType.isIntOrIndexOrFloat() && dstVectorType &&
2669	getElementTypeOrSelf(type: srcType) == getElementTypeOrSelf(dstVectorType))
2670	return BroadcastableToResult::Success;
2671	// From now on, only vectors broadcast.
2672	VectorType srcVectorType = llvm::dyn_cast<VectorType>(srcType);
2673	if (!srcVectorType)
2674	return BroadcastableToResult::SourceTypeNotAVector;
2675
2676	int64_t srcRank = srcVectorType.getRank();
2677	int64_t dstRank = dstVectorType.getRank();
2678	if (srcRank > dstRank)
2679	return BroadcastableToResult::SourceRankHigher;
2680	// Source has an exact match or singleton value for all trailing dimensions
2681	// (all leading dimensions are simply duplicated).
2682	int64_t lead = dstRank - srcRank;
2683	for (int64_t dimIdx = 0; dimIdx < srcRank; ++dimIdx) {
2684	// Have mismatching dims (in the sense of vector.broadcast semantics) been
2685	// encountered?
2686	bool foundMismatchingDims = false;
2687
2688	// Check fixed-width dims.
2689	int64_t srcDim = srcVectorType.getDimSize(dimIdx);
2690	int64_t dstDim = dstVectorType.getDimSize(lead + dimIdx);
2691	if (srcDim != 1 && srcDim != dstDim)
2692	foundMismatchingDims = true;
2693
2694	// Check scalable flags.
2695	bool srcDimScalableFlag = srcVectorType.getScalableDims()[dimIdx];
2696	bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + dimIdx];
2697	if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) ||
2698	// 1 -> [N] is fine, everything else should be rejected when mixing
2699	// fixed-width and scalable dims
2700	(srcDimScalableFlag != dstDimScalableFlag &&
2701	(srcDim != 1 || srcDimScalableFlag)))
2702	foundMismatchingDims = true;
2703
2704	if (foundMismatchingDims) {
2705	if (mismatchingDims != nullptr) {
2706	mismatchingDims->first.dim = srcDim;
2707	mismatchingDims->first.isScalable = srcDimScalableFlag;
2708
2709	mismatchingDims->second.dim = dstDim;
2710	mismatchingDims->second.isScalable = dstDimScalableFlag;
2711	}
2712	return BroadcastableToResult::DimensionMismatch;
2713	}
2714	}
2715
2716	return BroadcastableToResult::Success;
2717	}
2718
2719	LogicalResult BroadcastOp::verify() {
2720	std::pair<VectorDim, VectorDim> mismatchingDims;
2721	BroadcastableToResult res = isBroadcastableTo(
2722	getSourceType(), getResultVectorType(), &mismatchingDims);
2723	if (res == BroadcastableToResult::Success)
2724	return success();
2725	if (res == BroadcastableToResult::SourceRankHigher)
2726	return emitOpError("source rank higher than destination rank");
2727	if (res == BroadcastableToResult::DimensionMismatch) {
2728	return emitOpError("dimension mismatch (")
2729	<< (mismatchingDims.first.isScalable ? "[" : "")
2730	<< mismatchingDims.first.dim
2731	<< (mismatchingDims.first.isScalable ? "]" : "") << " vs. "
2732	<< (mismatchingDims.second.isScalable ? "[" : "")
2733	<< mismatchingDims.second.dim
2734	<< (mismatchingDims.second.isScalable ? "]" : "") << ")";
2735	}
2736	if (res == BroadcastableToResult::SourceTypeNotAVector)
2737	return emitOpError("source type is not a vector");
2738	llvm_unreachable("unexpected vector.broadcast op error");
2739	}
2740
2741	OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) {
2742	if (getSourceType() == getResultVectorType())
2743	return getSource();
2744	if (!adaptor.getSource())
2745	return {};
2746	auto vectorType = getResultVectorType();
2747	if (auto attr = llvm::dyn_cast<IntegerAttr>(adaptor.getSource())) {
2748	if (vectorType.getElementType() != attr.getType())
2749	return {};
2750	return DenseElementsAttr::get(vectorType, attr);
2751	}
2752	if (auto attr = llvm::dyn_cast<FloatAttr>(adaptor.getSource())) {
2753	if (vectorType.getElementType() != attr.getType())
2754	return {};
2755	return DenseElementsAttr::get(vectorType, attr);
2756	}
2757	if (auto attr = llvm::dyn_cast<SplatElementsAttr>(adaptor.getSource()))
2758	return DenseElementsAttr::get(vectorType, attr.getSplatValue<Attribute>());
2759	if (llvm::dyn_cast<ub::PoisonAttr>(adaptor.getSource()))
2760	return ub::PoisonAttr::get(getContext());
2761	return {};
2762	}
2763
2764	namespace {
2765
2766	// Fold broadcast1(broadcast2(x)) into broadcast1(x).
2767	struct BroadcastFolder : public OpRewritePattern<BroadcastOp> {
2768	using OpRewritePattern::OpRewritePattern;
2769
2770	LogicalResult matchAndRewrite(BroadcastOp broadcastOp,
2771	PatternRewriter &rewriter) const override {
2772	auto srcBroadcast = broadcastOp.getSource().getDefiningOp<BroadcastOp>();
2773	if (!srcBroadcast)
2774	return failure();
2775	rewriter.replaceOpWithNewOp<BroadcastOp>(broadcastOp,
2776	broadcastOp.getResultVectorType(),
2777	srcBroadcast.getSource());
2778	return success();
2779	}
2780	};
2781	} // namespace
2782
2783	void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
2784	MLIRContext *context) {
2785	// BroadcastToShapeCast is not a default canonicalization, it is opt-in by
2786	// calling `populateCastAwayVectorLeadingOneDimPatterns`
2787	results.add<BroadcastFolder>(context);
2788	}
2789
2790	//===----------------------------------------------------------------------===//
2791	// ShuffleOp
2792	//===----------------------------------------------------------------------===//
2793
2794	LogicalResult ShuffleOp::verify() {
2795	VectorType resultType = getResultVectorType();
2796	VectorType v1Type = getV1VectorType();
2797	VectorType v2Type = getV2VectorType();
2798	// Verify ranks.
2799	int64_t resRank = resultType.getRank();
2800	int64_t v1Rank = v1Type.getRank();
2801	int64_t v2Rank = v2Type.getRank();
2802	bool wellFormed0DCase = v1Rank == 0 && v2Rank == 0 && resRank == 1;
2803	bool wellFormedNDCase = v1Rank == resRank && v2Rank == resRank;
2804	if (!wellFormed0DCase && !wellFormedNDCase)
2805	return emitOpError("rank mismatch");
2806
2807	// Verify all but leading dimension sizes.
2808	for (int64_t r = 1; r < v1Rank; ++r) {
2809	int64_t resDim = resultType.getDimSize(r);
2810	int64_t v1Dim = v1Type.getDimSize(r);
2811	int64_t v2Dim = v2Type.getDimSize(r);
2812	if (resDim != v1Dim || v1Dim != v2Dim)
2813	return emitOpError("dimension mismatch");
2814	}
2815	// Verify mask length.
2816	ArrayRef<int64_t> mask = getMask();
2817	int64_t maskLength = mask.size();
2818	if (maskLength <= 0)
2819	return emitOpError("invalid mask length");
2820	if (maskLength != resultType.getDimSize(0))
2821	return emitOpError("mask length mismatch");
2822	// Verify all indices.
2823	int64_t indexSize = (v1Type.getRank() == 0 ? 1 : v1Type.getDimSize(0)) +
2824	(v2Type.getRank() == 0 ? 1 : v2Type.getDimSize(0));
2825	for (auto [idx, maskPos] : llvm::enumerate(mask)) {
2826	if (!isValidPositiveIndexOrPoison(maskPos, kPoisonIndex, indexSize))
2827	return emitOpError("mask index #") << (idx + 1) << " out of range";
2828	}
2829	return success();
2830	}
2831
2832	LogicalResult
2833	ShuffleOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
2834	ShuffleOp::Adaptor adaptor,
2835	SmallVectorImpl<Type> &inferredReturnTypes) {
2836	auto v1Type = llvm::cast<VectorType>(adaptor.getV1().getType());
2837	auto v1Rank = v1Type.getRank();
2838	// Construct resulting type: leading dimension matches mask
2839	// length, all trailing dimensions match the operands.
2840	SmallVector<int64_t, 4> shape;
2841	shape.reserve(v1Rank);
2842	shape.push_back(std::max<size_t>(1, adaptor.getMask().size()));
2843	// In the 0-D case there is no trailing shape to append.
2844	if (v1Rank > 0)
2845	llvm::append_range(shape, v1Type.getShape().drop_front());
2846	inferredReturnTypes.push_back(
2847	VectorType::get(shape, v1Type.getElementType()));
2848	return success();
2849	}
2850
2851	template <typename T>
2852	static bool isStepIndexArray(ArrayRef<T> idxArr, uint64_t begin, size_t width) {
2853	T expected = begin;
2854	return idxArr.size() == width && llvm::all_of(idxArr, [&expected](T value) {
2855	return value == expected++;
2856	});
2857	}
2858
2859	OpFoldResult vector::ShuffleOp::fold(FoldAdaptor adaptor) {
2860	auto v1Type = getV1VectorType();
2861	auto v2Type = getV2VectorType();
2862
2863	assert(!v1Type.isScalable() && !v2Type.isScalable() &&
2864	"Vector shuffle does not support scalable vectors");
2865
2866	// For consistency: 0-D shuffle return type is 1-D, this cannot be a folding
2867	// but must be a canonicalization into a vector.broadcast.
2868	if (v1Type.getRank() == 0)
2869	return {};
2870
2871	// Fold shuffle V1, V2, [0, 1, 2, 3] : <4xi32>, <2xi32> -> V1.
2872	auto mask = getMask();
2873	if (isStepIndexArray(mask, 0, v1Type.getDimSize(0)))
2874	return getV1();
2875	// Fold shuffle V1, V2, [4, 5] : <4xi32>, <2xi32> -> V2.
2876	if (isStepIndexArray(mask, v1Type.getDimSize(0), v2Type.getDimSize(0)))
2877	return getV2();
2878
2879	Attribute v1Attr = adaptor.getV1(), v2Attr = adaptor.getV2();
2880	if (!v1Attr || !v2Attr)
2881	return {};
2882
2883	// Fold shuffle poison, poison -> poison.
2884	bool isV1Poison = isa<ub::PoisonAttr>(v1Attr);
2885	bool isV2Poison = isa<ub::PoisonAttr>(v2Attr);
2886	if (isV1Poison && isV2Poison)
2887	return ub::PoisonAttr::get(getContext());
2888
2889	// Only support 1-D for now to avoid complicated n-D DenseElementsAttr
2890	// manipulation.
2891	if (v1Type.getRank() != 1)
2892	return {};
2893
2894	// Poison input attributes need special handling as they are not
2895	// DenseElementsAttr. If an index is poison, we select the first element of
2896	// the first non-poison input.
2897	SmallVector<Attribute> v1Elements, v2Elements;
2898	Attribute poisonElement;
2899	if (!isV2Poison) {
2900	v2Elements =
2901	to_vector(cast<DenseElementsAttr>(v2Attr).getValues<Attribute>());
2902	poisonElement = v2Elements[0];
2903	}
2904	if (!isV1Poison) {
2905	v1Elements =
2906	to_vector(cast<DenseElementsAttr>(v1Attr).getValues<Attribute>());
2907	poisonElement = v1Elements[0];
2908	}
2909
2910	SmallVector<Attribute> results;
2911	int64_t v1Size = v1Type.getDimSize(0);
2912	for (int64_t maskIdx : mask) {
2913	Attribute indexedElm;
2914	// TODO: Return a partial poison vector when supported by the UB dialect.
2915	if (maskIdx == ShuffleOp::kPoisonIndex) {
2916	indexedElm = poisonElement;
2917	} else {
2918	if (maskIdx < v1Size)
2919	indexedElm = isV1Poison ? poisonElement : v1Elements[maskIdx];
2920	else
2921	indexedElm = isV2Poison ? poisonElement : v2Elements[maskIdx - v1Size];
2922	}
2923
2924	results.push_back(indexedElm);
2925	}
2926
2927	return DenseElementsAttr::get(getResultVectorType(), results);
2928	}
2929
2930	namespace {
2931
2932	// Pattern to rewrite a 0-D shuffle with [0] or [1] mask returning a 1-D vector
2933	// to a broadcast.
2934	struct Canonicalize0DShuffleOp : public OpRewritePattern<ShuffleOp> {
2935	using OpRewritePattern::OpRewritePattern;
2936
2937	LogicalResult matchAndRewrite(ShuffleOp shuffleOp,
2938	PatternRewriter &rewriter) const override {
2939	VectorType v1VectorType = shuffleOp.getV1VectorType();
2940	ArrayRef<int64_t> mask = shuffleOp.getMask();
2941	if (v1VectorType.getRank() > 0)
2942	return failure();
2943	if (mask.size() != 1)
2944	return failure();
2945	VectorType resType = VectorType::Builder(v1VectorType).setShape({1});
2946	if (mask[0] == 0)
2947	rewriter.replaceOpWithNewOp<vector::BroadcastOp>(shuffleOp, resType,
2948	shuffleOp.getV1());
2949	else
2950	rewriter.replaceOpWithNewOp<vector::BroadcastOp>(shuffleOp, resType,
2951	shuffleOp.getV2());
2952	return success();
2953	}
2954	};
2955
2956	/// Pattern to rewrite a ShuffleOp(SplatOp, SplatOp) to SplatOp.
2957	class ShuffleSplat final : public OpRewritePattern<ShuffleOp> {
2958	public:
2959	using OpRewritePattern::OpRewritePattern;
2960
2961	LogicalResult matchAndRewrite(ShuffleOp op,
2962	PatternRewriter &rewriter) const override {
2963	auto v1Splat = op.getV1().getDefiningOp<SplatOp>();
2964	auto v2Splat = op.getV2().getDefiningOp<SplatOp>();
2965
2966	if (!v1Splat || !v2Splat)
2967	return failure();
2968
2969	if (v1Splat.getInput() != v2Splat.getInput())
2970	return failure();
2971
2972	rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), v1Splat.getInput());
2973	return success();
2974	}
2975	};
2976
2977	/// Pattern to rewrite a fixed-size interleave via vector.shuffle to
2978	/// vector.interleave.
2979	class ShuffleInterleave : public OpRewritePattern<ShuffleOp> {
2980	public:
2981	using OpRewritePattern::OpRewritePattern;
2982
2983	LogicalResult matchAndRewrite(ShuffleOp op,
2984	PatternRewriter &rewriter) const override {
2985	VectorType resultType = op.getResultVectorType();
2986	if (resultType.isScalable())
2987	return rewriter.notifyMatchFailure(
2988	op, "ShuffleOp can't represent a scalable interleave");
2989
2990	if (resultType.getRank() != 1)
2991	return rewriter.notifyMatchFailure(
2992	op, "ShuffleOp can't represent an n-D interleave");
2993
2994	VectorType sourceType = op.getV1VectorType();
2995	if (sourceType != op.getV2VectorType() ||
2996	sourceType.getNumElements() * 2 != resultType.getNumElements()) {
2997	return rewriter.notifyMatchFailure(
2998	op, "ShuffleOp types don't match an interleave");
2999	}
3000
3001	ArrayRef<int64_t> shuffleMask = op.getMask();
3002	int64_t resultVectorSize = resultType.getNumElements();
3003	for (int i = 0, e = resultVectorSize / 2; i < e; ++i) {
3004	int64_t maskValueA = shuffleMask[i * 2];
3005	int64_t maskValueB = shuffleMask[(i * 2) + 1];
3006	if (maskValueA != i || maskValueB != (resultVectorSize / 2) + i)
3007	return rewriter.notifyMatchFailure(op,
3008	"ShuffleOp mask not interleaving");
3009	}
3010
3011	rewriter.replaceOpWithNewOp<InterleaveOp>(op, op.getV1(), op.getV2());
3012	return success();
3013	}
3014	};
3015
3016	} // namespace
3017
3018	void ShuffleOp::getCanonicalizationPatterns(RewritePatternSet &results,
3019	MLIRContext *context) {
3020	results.add<ShuffleSplat, ShuffleInterleave, Canonicalize0DShuffleOp>(
3021	context);
3022	}
3023
3024	//===----------------------------------------------------------------------===//
3025	// InsertElementOp
3026	//===----------------------------------------------------------------------===//
3027
3028	void InsertElementOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
3029	SetIntRangeFn setResultRanges) {
3030	setResultRanges(getResult(), argRanges[0].rangeUnion(argRanges[1]));
3031	}
3032
3033	void InsertElementOp::build(OpBuilder &builder, OperationState &result,
3034	Value source, Value dest) {
3035	build(builder, result, source, dest, {});
3036	}
3037
3038	LogicalResult InsertElementOp::verify() {
3039	auto dstVectorType = getDestVectorType();
3040	if (dstVectorType.getRank() == 0) {
3041	if (getPosition())
3042	return emitOpError("expected position to be empty with 0-D vector");
3043	return success();
3044	}
3045	if (dstVectorType.getRank() != 1)
3046	return emitOpError("unexpected >1 vector rank");
3047	if (!getPosition())
3048	return emitOpError("expected position for 1-D vector");
3049	return success();
3050	}
3051
3052	OpFoldResult vector::InsertElementOp::fold(FoldAdaptor adaptor) {
3053	// Skip the 0-D vector here.
3054	if (!adaptor.getPosition())
3055	return {};
3056
3057	auto src = dyn_cast_or_null<TypedAttr>(adaptor.getSource());
3058	auto dst = dyn_cast_or_null<DenseElementsAttr>(adaptor.getDest());
3059	auto pos = dyn_cast_or_null<IntegerAttr>(adaptor.getPosition());
3060	if (!src || !dst || !pos)
3061	return {};
3062
3063	if (src.getType() != getDestVectorType().getElementType())
3064	return {};
3065
3066	auto dstElements = dst.getValues<Attribute>();
3067
3068	SmallVector<Attribute> results(dstElements);
3069
3070	uint64_t posIdx = pos.getInt();
3071	if (posIdx >= results.size())
3072	return {};
3073	results[posIdx] = src;
3074
3075	return DenseElementsAttr::get(getDestVectorType(), results);
3076	}
3077
3078	//===----------------------------------------------------------------------===//
3079	// InsertOp
3080	//===----------------------------------------------------------------------===//
3081
3082	void vector::InsertOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
3083	SetIntRangeFn setResultRanges) {
3084	setResultRanges(getResult(), argRanges[0].rangeUnion(argRanges[1]));
3085	}
3086
3087	void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
3088	Value source, Value dest) {
3089	auto vectorTy = cast<VectorType>(dest.getType());
3090	build(builder, result, source, dest,
3091	SmallVector<int64_t>(vectorTy.getRank(), 0));
3092	}
3093
3094	void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
3095	Value source, Value dest, int64_t position) {
3096	build(builder, result, source, dest, ArrayRef<int64_t>{position});
3097	}
3098
3099	void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
3100	Value source, Value dest, OpFoldResult position) {
3101	build(builder, result, source, dest, ArrayRef<OpFoldResult>{position});
3102	}
3103
3104	void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
3105	Value source, Value dest,
3106	ArrayRef<int64_t> position) {
3107	SmallVector<OpFoldResult> posVals;
3108	posVals.reserve(position.size());
3109	llvm::transform(position, std::back_inserter(posVals),
3110	[&](int64_t pos) { return builder.getI64IntegerAttr(pos); });
3111	build(builder, result, source, dest, posVals);
3112	}
3113
3114	void vector::InsertOp::build(OpBuilder &builder, OperationState &result,
3115	Value source, Value dest,
3116	ArrayRef<OpFoldResult> position) {
3117	SmallVector<int64_t> staticPos;
3118	SmallVector<Value> dynamicPos;
3119	dispatchIndexOpFoldResults(position, dynamicPos, staticPos);
3120	build(builder, result, source, dest, dynamicPos,
3121	builder.getDenseI64ArrayAttr(staticPos));
3122	}
3123
3124	LogicalResult InsertOp::verify() {
3125	SmallVector<OpFoldResult> position = getMixedPosition();
3126	auto destVectorType = getDestVectorType();
3127	if (position.size() > static_cast<unsigned>(destVectorType.getRank()))
3128	return emitOpError(
3129	"expected position attribute of rank no greater than dest vector rank");
3130	auto srcVectorType = llvm::dyn_cast<VectorType>(getValueToStoreType());
3131	if (srcVectorType &&
3132	(static_cast<unsigned>(srcVectorType.getRank()) + position.size() !=
3133	static_cast<unsigned>(destVectorType.getRank())))
3134	return emitOpError("expected position attribute rank + source rank to "
3135	"match dest vector rank");
3136	if (!srcVectorType &&
3137	(position.size() != static_cast<unsigned>(destVectorType.getRank())))
3138	return emitOpError(
3139	"expected position attribute rank to match the dest vector rank");
3140	for (auto [idx, pos] : llvm::enumerate(position)) {
3141	if (auto attr = dyn_cast<Attribute>(pos)) {
3142	int64_t constIdx = cast<IntegerAttr>(attr).getInt();
3143	if (!isValidPositiveIndexOrPoison(constIdx, kPoisonIndex,
3144	destVectorType.getDimSize(idx))) {
3145	return emitOpError("expected position attribute #")
3146	<< (idx + 1)
3147	<< " to be a non-negative integer smaller than the "
3148	"corresponding "
3149	"dest vector dimension";
3150	}
3151	}
3152	}
3153	return success();
3154	}
3155
3156	namespace {
3157
3158	// If insertOp is only inserting unit dimensions it can be transformed to a
3159	// broadcast.
3160	class InsertToBroadcast final : public OpRewritePattern<InsertOp> {
3161	public:
3162	using OpRewritePattern::OpRewritePattern;
3163
3164	LogicalResult matchAndRewrite(InsertOp insertOp,
3165	PatternRewriter &rewriter) const override {
3166	auto srcVecType =
3167	llvm::dyn_cast<VectorType>(insertOp.getValueToStoreType());
3168	if (!srcVecType || insertOp.getDestVectorType().getNumElements() !=
3169	srcVecType.getNumElements())
3170	return failure();
3171	rewriter.replaceOpWithNewOp<BroadcastOp>(
3172	insertOp, insertOp.getDestVectorType(), insertOp.getValueToStore());
3173	return success();
3174	}
3175	};
3176
3177	/// Pattern to rewrite a InsertOp(SplatOp, SplatOp) to SplatOp.
3178	class InsertSplatToSplat final : public OpRewritePattern<InsertOp> {
3179	public:
3180	using OpRewritePattern::OpRewritePattern;
3181
3182	LogicalResult matchAndRewrite(InsertOp op,
3183	PatternRewriter &rewriter) const override {
3184	auto srcSplat = op.getValueToStore().getDefiningOp<SplatOp>();
3185	auto dstSplat = op.getDest().getDefiningOp<SplatOp>();
3186
3187	if (!srcSplat || !dstSplat)
3188	return failure();
3189
3190	if (srcSplat.getInput() != dstSplat.getInput())
3191	return failure();
3192
3193	rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), srcSplat.getInput());
3194	return success();
3195	}
3196	};
3197
3198	} // namespace
3199
3200	static Attribute
3201	foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,
3202	Attribute dstAttr,
3203	int64_t maxVectorSizeFoldThreshold) {
3204	if (insertOp.hasDynamicPosition())
3205	return {};
3206
3207	auto denseDst = llvm::dyn_cast_if_present<DenseElementsAttr>(Val&: dstAttr);
3208	if (!denseDst)
3209	return {};
3210
3211	if (!srcAttr) {
3212	return {};
3213	}
3214
3215	VectorType destTy = insertOp.getDestVectorType();
3216	if (destTy.isScalable())
3217	return {};
3218
3219	// Make sure we do not create too many large constants.
3220	if (destTy.getNumElements() > maxVectorSizeFoldThreshold &&
3221	!insertOp->hasOneUse())
3222	return {};
3223
3224	// Calculate the linearized position of the continuous chunk of elements to
3225	// insert.
3226	llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0);
3227	copy(insertOp.getStaticPosition(), completePositions.begin());
3228	int64_t insertBeginPosition =
3229	linearize(completePositions, computeStrides(destTy.getShape()));
3230
3231	SmallVector<Attribute> insertedValues;
3232	Type destEltType = destTy.getElementType();
3233
3234	/// Converts the expected type to an IntegerAttr if there's
3235	/// a mismatch.
3236	auto convertIntegerAttr = [](Attribute attr, Type expectedType) -> Attribute {
3237	if (auto intAttr = mlir::dyn_cast<IntegerAttr>(attr)) {
3238	if (intAttr.getType() != expectedType)
3239	return IntegerAttr::get(expectedType, intAttr.getInt());
3240	}
3241	return attr;
3242	};
3243
3244	// The `convertIntegerAttr` method specifically handles the case
3245	// for `llvm.mlir.constant` which can hold an attribute with a
3246	// different type than the return type.
3247	if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(Val&: srcAttr)) {
3248	for (auto value : denseSource.getValues<Attribute>())
3249	insertedValues.push_back(convertIntegerAttr(value, destEltType));
3250	} else {
3251	insertedValues.push_back(Elt: convertIntegerAttr(srcAttr, destEltType));
3252	}
3253
3254	auto allValues = llvm::to_vector(denseDst.getValues<Attribute>());
3255	copy(insertedValues, allValues.begin() + insertBeginPosition);
3256	auto newAttr = DenseElementsAttr::get(destTy, allValues);
3257
3258	return newAttr;
3259	}
3260
3261	void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
3262	MLIRContext *context) {
3263	results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat>(context);
3264	}
3265
3266	OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
3267	// Do not create constants with more than `vectorSizeFoldThreashold` elements,
3268	// unless the source vector constant has a single use.
3269	constexpr int64_t vectorSizeFoldThreshold = 256;
3270	// Fold "vector.insert %v, %dest [] : vector<2x2xf32> from vector<2x2xf32>" to
3271	// %v. Note: Do not fold "vector.insert %v, %dest [] : f32 into vector<f32>"
3272	// (type mismatch).
3273	if (getNumIndices() == 0 && getValueToStoreType() == getType())
3274	return getValueToStore();
3275	// Fold `arith.constant` indices into the `vector.insert` operation. Make
3276	// sure that patterns requiring constant indices are added after this fold.
3277	SmallVector<Value> operands = {getValueToStore(), getDest()};
3278	if (auto val = extractInsertFoldConstantOp(*this, adaptor, operands))
3279	return val;
3280	if (auto res = foldPoisonIndexInsertExtractOp(
3281	getContext(), adaptor.getStaticPosition(), kPoisonIndex))
3282	return res;
3283	if (auto res = foldDenseElementsAttrDestInsertOp(
3284	*this, adaptor.getValueToStore(), adaptor.getDest(),
3285	vectorSizeFoldThreshold)) {
3286	return res;
3287	}
3288
3289	return {};
3290	}
3291
3292	//===----------------------------------------------------------------------===//
3293	// InsertStridedSliceOp
3294	//===----------------------------------------------------------------------===//
3295
3296	void InsertStridedSliceOp::build(OpBuilder &builder, OperationState &result,
3297	Value source, Value dest,
3298	ArrayRef<int64_t> offsets,
3299	ArrayRef<int64_t> strides) {
3300	result.addOperands({source, dest});
3301	auto offsetsAttr = getVectorSubscriptAttr(builder, offsets);
3302	auto stridesAttr = getVectorSubscriptAttr(builder, strides);
3303	result.addTypes(dest.getType());
3304	result.addAttribute(InsertStridedSliceOp::getOffsetsAttrName(result.name),
3305	offsetsAttr);
3306	result.addAttribute(InsertStridedSliceOp::getStridesAttrName(result.name),
3307	stridesAttr);
3308	}
3309
3310	// TODO: Should be moved to Tablegen ConfinedAttr attributes.
3311	template <typename OpType>
3312	static LogicalResult isIntegerArrayAttrSmallerThanShape(OpType op,
3313	ArrayAttr arrayAttr,
3314	ArrayRef<int64_t> shape,
3315	StringRef attrName) {
3316	if (arrayAttr.size() > shape.size())
3317	return op.emitOpError("expected ")
3318	<< attrName << " attribute of rank no greater than vector rank";
3319	return success();
3320	}
3321
3322	// Returns true if all integers in `arrayAttr` are in the half-open [min, max}
3323	// interval. If `halfOpen` is true then the admissible interval is [min, max).
3324	// Otherwise, the admissible interval is [min, max].
3325	template <typename OpType>
3326	static LogicalResult
3327	isIntegerArrayAttrConfinedToRange(OpType op, ArrayAttr arrayAttr, int64_t min,
3328	int64_t max, StringRef attrName,
3329	bool halfOpen = true) {
3330	for (auto attr : arrayAttr) {
3331	auto val = llvm::cast<IntegerAttr>(attr).getInt();
3332	auto upper = max;
3333	if (!halfOpen)
3334	upper += 1;
3335	if (val < min || val >= upper)
3336	return op.emitOpError("expected ") << attrName << " to be confined to ["
3337	<< min << ", " << upper << ")";
3338	}
3339	return success();
3340	}
3341
3342	// Returns true if all integers in `arrayAttr` are in the half-open [min, max}
3343	// interval. If `halfOpen` is true then the admissible interval is [min, max).
3344	// Otherwise, the admissible interval is [min, max].
3345	template <typename OpType>
3346	static LogicalResult
3347	isIntegerArrayAttrConfinedToShape(OpType op, ArrayAttr arrayAttr,
3348	ArrayRef<int64_t> shape, StringRef attrName,
3349	bool halfOpen = true, int64_t min = 0) {
3350	for (auto [index, attrDimPair] :
3351	llvm::enumerate(llvm::zip_first(arrayAttr, shape))) {
3352	int64_t val = llvm::cast<IntegerAttr>(std::get<0>(attrDimPair)).getInt();
3353	int64_t max = std::get<1>(attrDimPair);
3354	if (!halfOpen)
3355	max += 1;
3356	if (val < min || val >= max)
3357	return op.emitOpError("expected ")
3358	<< attrName << " dimension " << index << " to be confined to ["
3359	<< min << ", " << max << ")";
3360	}
3361	return success();
3362	}
3363
3364	// Returns true if, for all indices i = 0..shape.size()-1, val is in the
3365	// [min, max} interval:
3366	// val = `arrayAttr1[i]` + `arrayAttr2[i]`,
3367	// If `halfOpen` is true then the admissible interval is [min, max). Otherwise,
3368	// the admissible interval is [min, max].
3369	template <typename OpType>
3370	static LogicalResult isSumOfIntegerArrayAttrConfinedToShape(
3371	OpType op, ArrayAttr arrayAttr1, ArrayAttr arrayAttr2,
3372	ArrayRef<int64_t> shape, StringRef attrName1, StringRef attrName2,
3373	bool halfOpen = true, int64_t min = 1) {
3374	assert(arrayAttr1.size() <= shape.size());
3375	assert(arrayAttr2.size() <= shape.size());
3376	for (auto [index, it] :
3377	llvm::enumerate(llvm::zip(arrayAttr1, arrayAttr2, shape))) {
3378	auto val1 = llvm::cast<IntegerAttr>(std::get<0>(it)).getInt();
3379	auto val2 = llvm::cast<IntegerAttr>(std::get<1>(it)).getInt();
3380	int64_t max = std::get<2>(it);
3381	if (!halfOpen)
3382	max += 1;
3383	if (val1 + val2 < 0 || val1 + val2 >= max)
3384	return op.emitOpError("expected sum(")
3385	<< attrName1 << ", " << attrName2 << ") dimension " << index
3386	<< " to be confined to [" << min << ", " << max << ")";
3387	}
3388	return success();
3389	}
3390
3391	static ArrayAttr makeI64ArrayAttr(ArrayRef<int64_t> values,
3392	MLIRContext *context) {
3393	auto attrs = llvm::map_range(C&: values, F: [context](int64_t v) -> Attribute {
3394	return IntegerAttr::get(IntegerType::get(context, 64), APInt(64, v));
3395	});
3396	return ArrayAttr::get(context, llvm::to_vector<8>(attrs));
3397	}
3398
3399	LogicalResult InsertStridedSliceOp::verify() {
3400	auto sourceVectorType = getSourceVectorType();
3401	auto destVectorType = getDestVectorType();
3402	auto offsets = getOffsetsAttr();
3403	auto strides = getStridesAttr();
3404	if (offsets.size() != static_cast<unsigned>(destVectorType.getRank()))
3405	return emitOpError(
3406	"expected offsets of same size as destination vector rank");
3407	if (strides.size() != static_cast<unsigned>(sourceVectorType.getRank()))
3408	return emitOpError("expected strides of same size as source vector rank");
3409	if (sourceVectorType.getRank() > destVectorType.getRank())
3410	return emitOpError(
3411	"expected source rank to be no greater than destination rank");
3412
3413	auto sourceShape = sourceVectorType.getShape();
3414	auto destShape = destVectorType.getShape();
3415	SmallVector<int64_t, 4> sourceShapeAsDestShape(
3416	destShape.size() - sourceShape.size(), 0);
3417	sourceShapeAsDestShape.append(sourceShape.begin(), sourceShape.end());
3418	auto offName = InsertStridedSliceOp::getOffsetsAttrName();
3419	auto stridesName = InsertStridedSliceOp::getStridesAttrName();
3420	if (failed(isIntegerArrayAttrConfinedToShape(*this, offsets, destShape,
3421	offName)) ||
3422	failed(isIntegerArrayAttrConfinedToRange(*this, strides, /*min=*/1,
3423	/*max=*/1, stridesName,
3424	/*halfOpen=*/false)) ||
3425	failed(isSumOfIntegerArrayAttrConfinedToShape(
3426	*this, offsets,
3427	makeI64ArrayAttr(sourceShapeAsDestShape, getContext()), destShape,
3428	offName, "source vector shape",
3429	/*halfOpen=*/false, /*min=*/1)))
3430	return failure();
3431
3432	unsigned rankDiff = destShape.size() - sourceShape.size();
3433	for (unsigned idx = 0; idx < sourceShape.size(); ++idx) {
3434	if (sourceVectorType.getScalableDims()[idx] !=
3435	destVectorType.getScalableDims()[idx + rankDiff]) {
3436	return emitOpError("mismatching scalable flags (at source vector idx=")
3437	<< idx << ")";
3438	}
3439	if (sourceVectorType.getScalableDims()[idx]) {
3440	auto sourceSize = sourceShape[idx];
3441	auto destSize = destShape[idx + rankDiff];
3442	if (sourceSize != destSize) {
3443	return emitOpError("expected size at idx=")
3444	<< idx
3445	<< (" to match the corresponding base size from the input "
3446	"vector (")
3447	<< sourceSize << (" vs ") << destSize << (")");
3448	}
3449	}
3450	}
3451
3452	return success();
3453	}
3454
3455	namespace {
3456	/// Pattern to rewrite an InsertStridedSliceOp(SplatOp(X):src_type,
3457	/// SplatOp(X):dst_type) to SplatOp(X):dst_type.
3458	class FoldInsertStridedSliceSplat final
3459	: public OpRewritePattern<InsertStridedSliceOp> {
3460	public:
3461	using OpRewritePattern::OpRewritePattern;
3462
3463	LogicalResult matchAndRewrite(InsertStridedSliceOp insertStridedSliceOp,
3464	PatternRewriter &rewriter) const override {
3465	auto srcSplatOp =
3466	insertStridedSliceOp.getValueToStore().getDefiningOp<vector::SplatOp>();
3467	auto destSplatOp =
3468	insertStridedSliceOp.getDest().getDefiningOp<vector::SplatOp>();
3469
3470	if (!srcSplatOp || !destSplatOp)
3471	return failure();
3472
3473	if (srcSplatOp.getInput() != destSplatOp.getInput())
3474	return failure();
3475
3476	rewriter.replaceOp(insertStridedSliceOp, insertStridedSliceOp.getDest());
3477	return success();
3478	}
3479	};
3480
3481	/// Pattern to rewrite an InsertStridedSliceOp(ExtractStridedSliceOp(dst), dst)
3482	/// to dst.
3483	class FoldInsertStridedSliceOfExtract final
3484	: public OpRewritePattern<InsertStridedSliceOp> {
3485	public:
3486	using OpRewritePattern::OpRewritePattern;
3487
3488	LogicalResult matchAndRewrite(InsertStridedSliceOp insertStridedSliceOp,
3489	PatternRewriter &rewriter) const override {
3490	auto extractStridedSliceOp =
3491	insertStridedSliceOp.getValueToStore()
3492	.getDefiningOp<vector::ExtractStridedSliceOp>();
3493
3494	if (!extractStridedSliceOp)
3495	return failure();
3496
3497	if (extractStridedSliceOp.getOperand() != insertStridedSliceOp.getDest())
3498	return failure();
3499
3500	// Check if have the same strides and offsets.
3501	if (extractStridedSliceOp.getStrides() !=
3502	insertStridedSliceOp.getStrides() ||
3503	extractStridedSliceOp.getOffsets() != insertStridedSliceOp.getOffsets())
3504	return failure();
3505
3506	rewriter.replaceOp(insertStridedSliceOp, insertStridedSliceOp.getDest());
3507	return success();
3508	}
3509	};
3510
3511	// Pattern to rewrite an InsertStridedSliceOp(ConstantOp into ConstantOp) ->
3512	// ConstantOp.
3513	class InsertStridedSliceConstantFolder final
3514	: public OpRewritePattern<InsertStridedSliceOp> {
3515	public:
3516	using OpRewritePattern::OpRewritePattern;
3517
3518	// Do not create constants with more than `vectorSizeFoldThreashold` elements,
3519	// unless the source vector constant has a single use.
3520	static constexpr int64_t vectorSizeFoldThreshold = 256;
3521
3522	LogicalResult matchAndRewrite(InsertStridedSliceOp op,
3523	PatternRewriter &rewriter) const override {
3524	// Return if 'InsertOp' operand is not defined by a compatible vector
3525	// ConstantOp.
3526	TypedValue<VectorType> destVector = op.getDest();
3527	Attribute vectorDestCst;
3528	if (!matchPattern(value: destVector, pattern: m_Constant(bind_value: &vectorDestCst)))
3529	return failure();
3530
3531	VectorType destTy = destVector.getType();
3532	if (destTy.isScalable())
3533	return failure();
3534
3535	// Make sure we do not create too many large constants.
3536	if (destTy.getNumElements() > vectorSizeFoldThreshold &&
3537	!destVector.hasOneUse())
3538	return failure();
3539
3540	TypedValue<VectorType> sourceValue = op.getValueToStore();
3541	Attribute sourceCst;
3542	if (!matchPattern(value: sourceValue, pattern: m_Constant(bind_value: &sourceCst)))
3543	return failure();
3544
3545	// TODO: Support poison.
3546	if (isa<ub::PoisonAttr>(vectorDestCst) || isa<ub::PoisonAttr>(sourceCst))
3547	return failure();
3548
3549	// TODO: Handle non-unit strides when they become available.
3550	if (op.hasNonUnitStrides())
3551	return failure();
3552
3553	VectorType sliceVecTy = sourceValue.getType();
3554	ArrayRef<int64_t> sliceShape = sliceVecTy.getShape();
3555	int64_t rankDifference = destTy.getRank() - sliceVecTy.getRank();
3556	SmallVector<int64_t, 4> offsets = getI64SubArray(op.getOffsets());
3557	SmallVector<int64_t, 4> destStrides = computeStrides(destTy.getShape());
3558
3559	// Calcualte the destination element indices by enumerating all slice
3560	// positions within the destination and linearizing them. The enumeration
3561	// order is lexicographic which yields a sequence of monotonically
3562	// increasing linearized position indices.
3563	// Because the destination may have higher dimensionality then the slice,
3564	// we keep track of two overlapping sets of positions and offsets.
3565	auto denseDest = llvm::cast<DenseElementsAttr>(Val&: vectorDestCst);
3566	auto denseSlice = llvm::cast<DenseElementsAttr>(Val&: sourceCst);
3567	auto sliceValuesIt = denseSlice.value_begin<Attribute>();
3568	auto newValues = llvm::to_vector(denseDest.getValues<Attribute>());
3569	SmallVector<int64_t> currDestPosition(offsets.begin(), offsets.end());
3570	MutableArrayRef<int64_t> currSlicePosition(
3571	currDestPosition.begin() + rankDifference, currDestPosition.end());
3572	ArrayRef<int64_t> sliceOffsets(offsets.begin() + rankDifference,
3573	offsets.end());
3574	do {
3575	int64_t linearizedPosition = linearize(offsets: currDestPosition, basis: destStrides);
3576	assert(linearizedPosition < destTy.getNumElements() && "Invalid index");
3577	assert(sliceValuesIt != denseSlice.value_end<Attribute>() &&
3578	"Invalid slice element");
3579	newValues[linearizedPosition] = *sliceValuesIt;
3580	++sliceValuesIt;
3581	} while (succeeded(
3582	Result: incSlicePosition(position: currSlicePosition, shape: sliceShape, offsets: sliceOffsets)));
3583
3584	auto newAttr = DenseElementsAttr::get(destTy, newValues);
3585	rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
3586	return success();
3587	}
3588	};
3589
3590	} // namespace
3591
3592	void vector::InsertStridedSliceOp::getCanonicalizationPatterns(
3593	RewritePatternSet &results, MLIRContext *context) {
3594	results.add<FoldInsertStridedSliceSplat, FoldInsertStridedSliceOfExtract,
3595	InsertStridedSliceConstantFolder>(context);
3596	}
3597
3598	OpFoldResult InsertStridedSliceOp::fold(FoldAdaptor adaptor) {
3599	if (getSourceVectorType() == getDestVectorType())
3600	return getValueToStore();
3601	return {};
3602	}
3603
3604	//===----------------------------------------------------------------------===//
3605	// OuterProductOp
3606	//===----------------------------------------------------------------------===//
3607
3608	/// Build an op without mask, use the type of `acc` as the return type.
3609	void OuterProductOp::build(OpBuilder &builder, OperationState &result,
3610	Value lhs, Value rhs, Value acc) {
3611	result.addOperands({lhs, rhs, acc});
3612	result.addTypes(acc.getType());
3613	}
3614
3615	void OuterProductOp::print(OpAsmPrinter &p) {
3616	p << " " << getLhs() << ", " << getRhs();
3617	if (getAcc()) {
3618	p << ", " << getAcc();
3619	p.printOptionalAttrDict((*this)->getAttrs());
3620	}
3621	p << " : " << getLhs().getType() << ", " << getRhs().getType();
3622	}
3623
3624	ParseResult OuterProductOp::parse(OpAsmParser &parser, OperationState &result) {
3625	SmallVector<OpAsmParser::UnresolvedOperand, 3> operandsInfo;
3626	Type tLHS, tRHS;
3627	if (parser.parseOperandList(operandsInfo) ||
3628	parser.parseOptionalAttrDict(result.attributes) ||
3629	parser.parseColonType(tLHS) || parser.parseComma() ||
3630	parser.parseType(tRHS))
3631	return failure();
3632	if (operandsInfo.size() < 2)
3633	return parser.emitError(parser.getNameLoc(),
3634	"expected at least 2 operands");
3635	VectorType vLHS = llvm::dyn_cast<VectorType>(tLHS);
3636	VectorType vRHS = llvm::dyn_cast<VectorType>(tRHS);
3637	if (!vLHS)
3638	return parser.emitError(parser.getNameLoc(),
3639	"expected vector type for operand #1");
3640
3641	VectorType resType;
3642	if (vRHS) {
3643	SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0],
3644	vRHS.getScalableDims()[0]};
3645	resType = VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
3646	vLHS.getElementType(), scalableDimsRes);
3647	} else {
3648	// Scalar RHS operand
3649	SmallVector<bool> scalableDimsRes{vLHS.getScalableDims()[0]};
3650	resType = VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType(),
3651	scalableDimsRes);
3652	}
3653
3654	if (!result.attributes.get(OuterProductOp::getKindAttrName(result.name))) {
3655	result.attributes.append(
3656	OuterProductOp::getKindAttrName(result.name),
3657	CombiningKindAttr::get(result.getContext(),
3658	OuterProductOp::getDefaultKind()));
3659	}
3660
3661	return failure(
3662	parser.resolveOperand(operandsInfo[0], tLHS, result.operands) ||
3663	parser.resolveOperand(operandsInfo[1], tRHS, result.operands) ||
3664	(operandsInfo.size() > 2 &&
3665	parser.resolveOperand(operandsInfo[2], resType, result.operands)) ||
3666	parser.addTypeToList(resType, result.types));
3667	}
3668
3669	LogicalResult OuterProductOp::verify() {
3670	Type tRHS = getOperandTypeRHS();
3671	VectorType vLHS = getOperandVectorTypeLHS(),
3672	vRHS = llvm::dyn_cast<VectorType>(tRHS),
3673	vACC = getOperandVectorTypeACC(), vRES = getResultVectorType();
3674
3675	if (vLHS.getRank() != 1)
3676	return emitOpError("expected 1-d vector for operand #1");
3677
3678	if (vRHS) {
3679	// Proper OUTER operation.
3680	if (vRHS.getRank() != 1)
3681	return emitOpError("expected 1-d vector for operand #2");
3682	if (vRES.getRank() != 2)
3683	return emitOpError("expected 2-d vector result");
3684	if (vLHS.getDimSize(0) != vRES.getDimSize(0))
3685	return emitOpError("expected #1 operand dim to match result dim #1");
3686	if (vRHS.getDimSize(0) != vRES.getDimSize(1))
3687	return emitOpError("expected #2 operand dim to match result dim #2");
3688	if (vLHS.isScalable() && !vRHS.isScalable()) {
3689	// This restriction reflects what's currently supported in terms of
3690	// scalable vectors. However, we could relax this if there's a use case.
3691	return emitOpError(
3692	"expected either both or only #2 operand dim to be scalable");
3693	}
3694	} else {
3695	// An AXPY operation.
3696	if (vRES.getRank() != 1)
3697	return emitOpError("expected 1-d vector result");
3698	if (vLHS.getDimSize(0) != vRES.getDimSize(0))
3699	return emitOpError("expected #1 operand dim to match result dim #1");
3700	}
3701
3702	if (vACC && vACC != vRES)
3703	return emitOpError("expected operand #3 of same type as result type");
3704
3705	// Verify supported combining kind.
3706	if (!isSupportedCombiningKind(getKind(), vRES.getElementType()))
3707	return emitOpError("unsupported outerproduct type");
3708
3709	return success();
3710	}
3711
3712	// MaskableOpInterface methods.
3713
3714	/// Returns the mask type expected by this operation. Mostly used for
3715	/// verification purposes. It requires the operation to be vectorized."
3716	Type OuterProductOp::getExpectedMaskType() {
3717	auto vecType = this->getResultVectorType();
3718	return VectorType::get(vecType.getShape(),
3719	IntegerType::get(vecType.getContext(), /*width=*/1),
3720	vecType.getScalableDims());
3721	}
3722
3723	//===----------------------------------------------------------------------===//
3724	// ExtractStridedSliceOp
3725	//===----------------------------------------------------------------------===//
3726
3727	// Inference works as follows:
3728	// 1. Add 'sizes' from prefix of dims in 'offsets'.
3729	// 2. Add sizes from 'vectorType' for remaining dims.
3730	// Scalable flags are inherited from 'vectorType'.
3731	static Type inferStridedSliceOpResultType(VectorType vectorType,
3732	ArrayAttr offsets, ArrayAttr sizes,
3733	ArrayAttr strides) {
3734	assert(offsets.size() == sizes.size() && offsets.size() == strides.size());
3735	SmallVector<int64_t, 4> shape;
3736	shape.reserve(N: vectorType.getRank());
3737	unsigned idx = 0;
3738	for (unsigned e = offsets.size(); idx < e; ++idx)
3739	shape.push_back(Elt: llvm::cast<IntegerAttr>(sizes[idx]).getInt());
3740	for (unsigned e = vectorType.getShape().size(); idx < e; ++idx)
3741	shape.push_back(Elt: vectorType.getShape()[idx]);
3742
3743	return VectorType::get(shape, vectorType.getElementType(),
3744	vectorType.getScalableDims());
3745	}
3746
3747	void ExtractStridedSliceOp::build(OpBuilder &builder, OperationState &result,
3748	Value source, ArrayRef<int64_t> offsets,
3749	ArrayRef<int64_t> sizes,
3750	ArrayRef<int64_t> strides) {
3751	result.addOperands(source);
3752	auto offsetsAttr = getVectorSubscriptAttr(builder, offsets);
3753	auto sizesAttr = getVectorSubscriptAttr(builder, sizes);
3754	auto stridesAttr = getVectorSubscriptAttr(builder, strides);
3755	result.addTypes(
3756	inferStridedSliceOpResultType(llvm::cast<VectorType>(source.getType()),
3757	offsetsAttr, sizesAttr, stridesAttr));
3758	result.addAttribute(ExtractStridedSliceOp::getOffsetsAttrName(result.name),
3759	offsetsAttr);
3760	result.addAttribute(ExtractStridedSliceOp::getSizesAttrName(result.name),
3761	sizesAttr);
3762	result.addAttribute(ExtractStridedSliceOp::getStridesAttrName(result.name),
3763	stridesAttr);
3764	}
3765
3766	LogicalResult ExtractStridedSliceOp::verify() {
3767	auto type = getSourceVectorType();
3768	auto offsets = getOffsetsAttr();
3769	auto sizes = getSizesAttr();
3770	auto strides = getStridesAttr();
3771	if (offsets.size() != sizes.size() || offsets.size() != strides.size())
3772	return emitOpError(
3773	"expected offsets, sizes and strides attributes of same size");
3774
3775	auto shape = type.getShape();
3776	auto offName = getOffsetsAttrName();
3777	auto sizesName = getSizesAttrName();
3778	auto stridesName = getStridesAttrName();
3779	if (failed(
3780	isIntegerArrayAttrSmallerThanShape(*this, offsets, shape, offName)) ||
3781	failed(
3782	isIntegerArrayAttrSmallerThanShape(*this, sizes, shape, sizesName)) ||
3783	failed(isIntegerArrayAttrSmallerThanShape(*this, strides, shape,
3784	stridesName)) ||
3785	failed(
3786	isIntegerArrayAttrConfinedToShape(*this, offsets, shape, offName)) ||
3787	failed(isIntegerArrayAttrConfinedToShape(*this, sizes, shape, sizesName,
3788	/*halfOpen=*/false,
3789	/*min=*/1)) ||
3790	failed(isIntegerArrayAttrConfinedToRange(*this, strides, /*min=*/1,
3791	/*max=*/1, stridesName,
3792	/*halfOpen=*/false)) ||
3793	failed(isSumOfIntegerArrayAttrConfinedToShape(*this, offsets, sizes,
3794	shape, offName, sizesName,
3795	/*halfOpen=*/false)))
3796	return failure();
3797
3798	auto resultType = inferStridedSliceOpResultType(getSourceVectorType(),
3799	offsets, sizes, strides);
3800	if (getResult().getType() != resultType)
3801	return emitOpError("expected result type to be ") << resultType;
3802
3803	for (unsigned idx = 0; idx < sizes.size(); ++idx) {
3804	if (type.getScalableDims()[idx]) {
3805	auto inputDim = type.getShape()[idx];
3806	auto inputSize = llvm::cast<IntegerAttr>(sizes[idx]).getInt();
3807	if (inputDim != inputSize)
3808	return emitOpError("expected size at idx=")
3809	<< idx
3810	<< (" to match the corresponding base size from the input "
3811	"vector (")
3812	<< inputSize << (" vs ") << inputDim << (")");
3813	}
3814	}
3815
3816	return success();
3817	}
3818
3819	// When the source of ExtractStrided comes from a chain of InsertStrided ops try
3820	// to use the source of the InsertStrided ops if we can detect that the
3821	// extracted vector is a subset of one of the vector inserted.
3822	static LogicalResult
3823	foldExtractStridedOpFromInsertChain(ExtractStridedSliceOp op) {
3824	// Helper to extract integer out of ArrayAttr.
3825	auto getElement = [](ArrayAttr array, int idx) {
3826	return llvm::cast<IntegerAttr>(array[idx]).getInt();
3827	};
3828	ArrayAttr extractOffsets = op.getOffsets();
3829	ArrayAttr extractStrides = op.getStrides();
3830	ArrayAttr extractSizes = op.getSizes();
3831	auto insertOp = op.getVector().getDefiningOp<InsertStridedSliceOp>();
3832	while (insertOp) {
3833	if (op.getSourceVectorType().getRank() !=
3834	insertOp.getSourceVectorType().getRank())
3835	return failure();
3836	ArrayAttr insertOffsets = insertOp.getOffsets();
3837	ArrayAttr insertStrides = insertOp.getStrides();
3838	// If the rank of extract is greater than the rank of insert, we are likely
3839	// extracting a partial chunk of the vector inserted.
3840	if (extractOffsets.size() > insertOffsets.size())
3841	return failure();
3842	bool patialoverlap = false;
3843	bool disjoint = false;
3844	SmallVector<int64_t, 4> offsetDiffs;
3845	for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
3846	if (getElement(extractStrides, dim) != getElement(insertStrides, dim))
3847	return failure();
3848	int64_t start = getElement(insertOffsets, dim);
3849	int64_t end = start + insertOp.getSourceVectorType().getDimSize(dim);
3850	int64_t offset = getElement(extractOffsets, dim);
3851	int64_t size = getElement(extractSizes, dim);
3852	// Check if the start of the extract offset is in the interval inserted.
3853	if (start <= offset && offset < end) {
3854	// If the extract interval overlaps but is not fully included we may
3855	// have a partial overlap that will prevent any folding.
3856	if (offset + size > end)
3857	patialoverlap = true;
3858	offsetDiffs.push_back(Elt: offset - start);
3859	continue;
3860	}
3861	disjoint = true;
3862	break;
3863	}
3864	// The extract element chunk is a subset of the insert element.
3865	if (!disjoint && !patialoverlap) {
3866	op.setOperand(insertOp.getValueToStore());
3867	// OpBuilder is only used as a helper to build an I64ArrayAttr.
3868	OpBuilder b(op.getContext());
3869	op.setOffsetsAttr(b.getI64ArrayAttr(offsetDiffs));
3870	return success();
3871	}
3872	// If the chunk extracted is disjoint from the chunk inserted, keep looking
3873	// in the insert chain.
3874	if (disjoint)
3875	insertOp = insertOp.getDest().getDefiningOp<InsertStridedSliceOp>();
3876	else {
3877	// The extracted vector partially overlap the inserted vector, we cannot
3878	// fold.
3879	return failure();
3880	}
3881	}
3882	return failure();
3883	}
3884
3885	// ExtractStridedSliceOp(non-splat ConstantOp) -> ConstantOp.
3886	static OpFoldResult
3887	foldExtractStridedSliceNonSplatConstant(ExtractStridedSliceOp op,
3888	Attribute foldInput) {
3889
3890	auto dense = llvm::dyn_cast_if_present<DenseElementsAttr>(Val&: foldInput);
3891	if (!dense)
3892	return {};
3893
3894	// TODO: Handle non-unit strides when they become available.
3895	if (op.hasNonUnitStrides())
3896	return {};
3897
3898	VectorType sourceVecTy = op.getSourceVectorType();
3899	ArrayRef<int64_t> sourceShape = sourceVecTy.getShape();
3900	SmallVector<int64_t, 4> sourceStrides = computeStrides(sizes: sourceShape);
3901
3902	VectorType sliceVecTy = op.getType();
3903	ArrayRef<int64_t> sliceShape = sliceVecTy.getShape();
3904	int64_t rank = sliceVecTy.getRank();
3905
3906	// Expand offsets and sizes to match the vector rank.
3907	SmallVector<int64_t, 4> offsets(rank, 0);
3908	copy(getI64SubArray(op.getOffsets()), offsets.begin());
3909
3910	SmallVector<int64_t, 4> sizes(sourceShape);
3911	copy(getI64SubArray(op.getSizes()), sizes.begin());
3912
3913	// Calculate the slice elements by enumerating all slice positions and
3914	// linearizing them. The enumeration order is lexicographic which yields a
3915	// sequence of monotonically increasing linearized position indices.
3916	const auto denseValuesBegin = dense.value_begin<Attribute>();
3917	SmallVector<Attribute> sliceValues;
3918	sliceValues.reserve(N: sliceVecTy.getNumElements());
3919	SmallVector<int64_t> currSlicePosition(offsets.begin(), offsets.end());
3920	do {
3921	int64_t linearizedPosition = linearize(offsets: currSlicePosition, basis: sourceStrides);
3922	assert(linearizedPosition < sourceVecTy.getNumElements() &&
3923	"Invalid index");
3924	sliceValues.push_back(Elt: *(denseValuesBegin + linearizedPosition));
3925	} while (succeeded(Result: incSlicePosition(position: currSlicePosition, shape: sliceShape, offsets)));
3926
3927	assert(static_cast<int64_t>(sliceValues.size()) ==
3928	sliceVecTy.getNumElements() &&
3929	"Invalid number of slice elements");
3930	return DenseElementsAttr::get(sliceVecTy, sliceValues);
3931	}
3932
3933	OpFoldResult ExtractStridedSliceOp::fold(FoldAdaptor adaptor) {
3934	if (getSourceVectorType() == getResult().getType())
3935	return getVector();
3936	if (succeeded(foldExtractStridedOpFromInsertChain(*this)))
3937	return getResult();
3938
3939	// ExtractStridedSliceOp(splat ConstantOp) -> ConstantOp.
3940	if (auto splat =
3941	llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getVector()))
3942	DenseElementsAttr::get(getType(), splat.getSplatValue<Attribute>());
3943
3944	// ExtractStridedSliceOp(non-splat ConstantOp) -> ConstantOp.
3945	return foldExtractStridedSliceNonSplatConstant(*this, adaptor.getVector());
3946	}
3947
3948	void ExtractStridedSliceOp::getOffsets(SmallVectorImpl<int64_t> &results) {
3949	populateFromInt64AttrArray(getOffsets(), results);
3950	}
3951
3952	namespace {
3953
3954	// Pattern to rewrite an ExtractStridedSliceOp(ConstantMaskOp) to
3955	// ConstantMaskOp.
3956	class StridedSliceConstantMaskFolder final
3957	: public OpRewritePattern<ExtractStridedSliceOp> {
3958	public:
3959	using OpRewritePattern::OpRewritePattern;
3960
3961	LogicalResult matchAndRewrite(ExtractStridedSliceOp extractStridedSliceOp,
3962	PatternRewriter &rewriter) const override {
3963	// Return if 'extractStridedSliceOp' operand is not defined by a
3964	// ConstantMaskOp.
3965	auto *defOp = extractStridedSliceOp.getVector().getDefiningOp();
3966	auto constantMaskOp = dyn_cast_or_null<ConstantMaskOp>(defOp);
3967	if (!constantMaskOp)
3968	return failure();
3969	// Return if 'extractStridedSliceOp' has non-unit strides.
3970	if (extractStridedSliceOp.hasNonUnitStrides())
3971	return failure();
3972	// Gather constant mask dimension sizes.
3973	ArrayRef<int64_t> maskDimSizes = constantMaskOp.getMaskDimSizes();
3974	// Gather strided slice offsets and sizes.
3975	SmallVector<int64_t, 4> sliceOffsets;
3976	populateFromInt64AttrArray(extractStridedSliceOp.getOffsets(),
3977	sliceOffsets);
3978	SmallVector<int64_t, 4> sliceSizes;
3979	populateFromInt64AttrArray(extractStridedSliceOp.getSizes(), sliceSizes);
3980
3981	// Compute slice of vector mask region.
3982	SmallVector<int64_t, 4> sliceMaskDimSizes;
3983	sliceMaskDimSizes.reserve(N: maskDimSizes.size());
3984	for (auto [maskDimSize, sliceOffset, sliceSize] :
3985	llvm::zip(maskDimSizes, sliceOffsets, sliceSizes)) {
3986	int64_t sliceMaskDimSize = std::max(
3987	static_cast<int64_t>(0),
3988	std::min(sliceOffset + sliceSize, maskDimSize) - sliceOffset);
3989	sliceMaskDimSizes.push_back(sliceMaskDimSize);
3990	}
3991	// Add unchanged dimensions.
3992	if (sliceMaskDimSizes.size() < maskDimSizes.size())
3993	for (size_t i = sliceMaskDimSizes.size(); i < maskDimSizes.size(); ++i)
3994	sliceMaskDimSizes.push_back(Elt: maskDimSizes[i]);
3995	// If any of 'sliceMaskDimSizes' are zero, then set all to zero (masked
3996	// region is a conjunction of mask dim intervals).
3997	if (llvm::is_contained(Range&: sliceMaskDimSizes, Element: 0))
3998	sliceMaskDimSizes.assign(NumElts: maskDimSizes.size(), Elt: 0);
3999
4000	// Replace 'extractStridedSliceOp' with ConstantMaskOp with sliced mask
4001	// region.
4002	rewriter.replaceOpWithNewOp<ConstantMaskOp>(
4003	extractStridedSliceOp, extractStridedSliceOp.getResult().getType(),
4004	sliceMaskDimSizes);
4005	return success();
4006	}
4007	};
4008
4009	// Pattern to rewrite an ExtractStridedSliceOp(BroadcastOp) to
4010	// BroadcastOp(ExtractStrideSliceOp).
4011	class StridedSliceBroadcast final
4012	: public OpRewritePattern<ExtractStridedSliceOp> {
4013	public:
4014	using OpRewritePattern::OpRewritePattern;
4015
4016	LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
4017	PatternRewriter &rewriter) const override {
4018	auto broadcast = op.getVector().getDefiningOp<BroadcastOp>();
4019	if (!broadcast)
4020	return failure();
4021	auto srcVecType =
4022	llvm::dyn_cast<VectorType>(broadcast.getSource().getType());
4023	unsigned srcRank = srcVecType ? srcVecType.getRank() : 0;
4024	auto dstVecType = llvm::cast<VectorType>(op.getType());
4025	unsigned dstRank = dstVecType.getRank();
4026	unsigned rankDiff = dstRank - srcRank;
4027	// Check if the most inner dimensions of the source of the broadcast are the
4028	// same as the destination of the extract. If this is the case we can just
4029	// use a broadcast as the original dimensions are untouched.
4030	bool lowerDimMatch = true;
4031	for (unsigned i = 0; i < srcRank; i++) {
4032	if (srcVecType.getDimSize(i) != dstVecType.getDimSize(i + rankDiff)) {
4033	lowerDimMatch = false;
4034	break;
4035	}
4036	}
4037	Value source = broadcast.getSource();
4038	// If the inner dimensions don't match, it means we need to extract from the
4039	// source of the orignal broadcast and then broadcast the extracted value.
4040	// We also need to handle degenerated cases where the source is effectively
4041	// just a single scalar.
4042	bool isScalarSrc = (srcRank == 0 || srcVecType.getNumElements() == 1);
4043	if (!lowerDimMatch && !isScalarSrc) {
4044	source = rewriter.create<ExtractStridedSliceOp>(
4045	op->getLoc(), source,
4046	getI64SubArray(op.getOffsets(), /* dropFront=*/rankDiff),
4047	getI64SubArray(op.getSizes(), /* dropFront=*/rankDiff),
4048	getI64SubArray(op.getStrides(), /* dropFront=*/rankDiff));
4049	}
4050	rewriter.replaceOpWithNewOp<BroadcastOp>(op, op.getType(), source);
4051	return success();
4052	}
4053	};
4054
4055	/// Pattern to rewrite an ExtractStridedSliceOp(SplatOp) to SplatOp.
4056	class StridedSliceSplat final : public OpRewritePattern<ExtractStridedSliceOp> {
4057	public:
4058	using OpRewritePattern::OpRewritePattern;
4059
4060	LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
4061	PatternRewriter &rewriter) const override {
4062	auto splat = op.getVector().getDefiningOp<SplatOp>();
4063	if (!splat)
4064	return failure();
4065	rewriter.replaceOpWithNewOp<SplatOp>(op, op.getType(), splat.getInput());
4066	return success();
4067	}
4068	};
4069
4070	/// Pattern to rewrite simple cases of N-D extract_strided_slice, where the
4071	/// slice is contiguous, into extract and shape_cast.
4072	///
4073	/// Example:
4074	/// Before:
4075	/// %1 = vector.extract_strided_slice %arg0 {
4076	/// offsets = [0, 0, 0, 0, 0],
4077	/// sizes = [1, 1, 1, 1, 8],
4078	/// strides = [1, 1, 1, 1, 1]
4079	/// } : vector<8x1x1x2x8xi8> to vector<1x1x1x1x8xi8>
4080	/// After:
4081	/// %0 = vector.extract %arg0[0, 0, 0, 0]
4082	/// : vector<8xi8> from vector<8x1x1x2x8xi8>
4083	/// %1 = vector.shape_cast %0
4084	/// : vector<8xi8> to vector<1x1x1x1x8xi8>
4085	///
4086	class ContiguousExtractStridedSliceToExtract final
4087	: public OpRewritePattern<ExtractStridedSliceOp> {
4088	public:
4089	using OpRewritePattern::OpRewritePattern;
4090
4091	LogicalResult matchAndRewrite(ExtractStridedSliceOp op,
4092	PatternRewriter &rewriter) const override {
4093	if (op.hasNonUnitStrides())
4094	return failure();
4095	Value source = op.getOperand();
4096	auto sourceType = cast<VectorType>(source.getType());
4097	if (sourceType.isScalable() || sourceType.getRank() == 0)
4098	return failure();
4099
4100	// Compute the number of offsets to pass to ExtractOp::build. That is the
4101	// difference between the source rank and the desired slice rank. We walk
4102	// the dimensions from innermost out, and stop when the next slice dimension
4103	// is not full-size.
4104	SmallVector<int64_t> sizes = getI64SubArray(op.getSizes());
4105	int numOffsets;
4106	for (numOffsets = sizes.size(); numOffsets > 0; --numOffsets) {
4107	if (sizes[numOffsets - 1] != sourceType.getDimSize(numOffsets - 1))
4108	break;
4109	}
4110
4111	// If the created extract op would have no offsets, then this whole
4112	// extract_strided_slice is the identity and should have been handled by
4113	// other canonicalizations.
4114	if (numOffsets == 0)
4115	return failure();
4116
4117	// If not even the inner-most dimension is full-size, this op can't be
4118	// rewritten as an ExtractOp.
4119	if (numOffsets == sourceType.getRank() &&
4120	static_cast<int>(sizes.size()) == sourceType.getRank())
4121	return failure();
4122
4123	// The outer dimensions must have unit size.
4124	for (int i = 0; i < numOffsets; ++i) {
4125	if (sizes[i] != 1)
4126	return failure();
4127	}
4128
4129	// Avoid generating slices that have leading unit dimensions. The shape_cast
4130	// op that we create below would take bad generic fallback patterns
4131	// (ShapeCastOpRewritePattern).
4132	while (numOffsets < static_cast<int>(sizes.size()) - 1 &&
4133	sizes[numOffsets] == 1) {
4134	++numOffsets;
4135	}
4136
4137	SmallVector<int64_t> offsets = getI64SubArray(op.getOffsets());
4138	auto extractOffsets = ArrayRef(offsets).take_front(N: numOffsets);
4139	Value extract = rewriter.create<vector::ExtractOp>(op->getLoc(), source,
4140	extractOffsets);
4141	rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getType(), extract);
4142	return success();
4143	}
4144	};
4145
4146	} // namespace
4147
4148	void ExtractStridedSliceOp::getCanonicalizationPatterns(
4149	RewritePatternSet &results, MLIRContext *context) {
4150	// Pattern to rewrite a ExtractStridedSliceOp(ConstantMaskOp) ->
4151	// ConstantMaskOp and ExtractStridedSliceOp(ConstantOp) -> ConstantOp.
4152	results.add<StridedSliceConstantMaskFolder, StridedSliceBroadcast,
4153	StridedSliceSplat, ContiguousExtractStridedSliceToExtract>(
4154	context);
4155	}
4156
4157	//===----------------------------------------------------------------------===//
4158	// TransferReadOp
4159	//===----------------------------------------------------------------------===//
4160
4161	/// 1. Builder that sets padding to zero and an empty mask (variant with attrs).
4162	void TransferReadOp::build(OpBuilder &builder, OperationState &result,
4163	VectorType vectorType, Value source,
4164	ValueRange indices, AffineMapAttr permutationMapAttr,
4165	/*optional*/ ArrayAttr inBoundsAttr) {
4166	Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
4167	Value padding = builder.create<arith::ConstantOp>(
4168	result.location, elemType, builder.getZeroAttr(elemType));
4169	build(builder, result, vectorType, source, indices, permutationMapAttr,
4170	padding, /*mask=*/Value(), inBoundsAttr);
4171	}
4172
4173	/// 2. Builder that sets padding to zero an empty mask (variant without attrs).
4174	void TransferReadOp::build(OpBuilder &builder, OperationState &result,
4175	VectorType vectorType, Value source,
4176	ValueRange indices, AffineMap permutationMap,
4177	std::optional<ArrayRef<bool>> inBounds) {
4178	auto permutationMapAttr = AffineMapAttr::get(permutationMap);
4179	auto inBoundsAttr = (inBounds && !inBounds.value().empty())
4180	? builder.getBoolArrayAttr(inBounds.value())
4181	: builder.getBoolArrayAttr(
4182	SmallVector<bool>(vectorType.getRank(), false));
4183	build(builder, result, vectorType, source, indices, permutationMapAttr,
4184	inBoundsAttr);
4185	}
4186
4187	/// 3. Builder that sets permutation map to 'getMinorIdentityMap'.
4188	void TransferReadOp::build(OpBuilder &builder, OperationState &result,
4189	VectorType vectorType, Value source,
4190	ValueRange indices, Value padding,
4191	std::optional<ArrayRef<bool>> inBounds) {
4192	AffineMap permutationMap = getTransferMinorIdentityMap(
4193	llvm::cast<ShapedType>(source.getType()), vectorType);
4194	auto permutationMapAttr = AffineMapAttr::get(permutationMap);
4195	auto inBoundsAttr = (inBounds && !inBounds.value().empty())
4196	? builder.getBoolArrayAttr(inBounds.value())
4197	: builder.getBoolArrayAttr(
4198	SmallVector<bool>(vectorType.getRank(), false));
4199	build(builder, result, vectorType, source, indices, permutationMapAttr,
4200	padding,
4201	/*mask=*/Value(), inBoundsAttr);
4202	}
4203
4204	/// 4. Builder that sets padding to zero and permutation map to
4205	/// 'getMinorIdentityMap'.
4206	void TransferReadOp::build(OpBuilder &builder, OperationState &result,
4207	VectorType vectorType, Value source,
4208	ValueRange indices,
4209	std::optional<ArrayRef<bool>> inBounds) {
4210	Type elemType = llvm::cast<ShapedType>(source.getType()).getElementType();
4211	Value padding = builder.create<arith::ConstantOp>(
4212	result.location, elemType, builder.getZeroAttr(elemType));
4213	build(builder, result, vectorType, source, indices, padding, inBounds);
4214	}
4215
4216	template <typename EmitFun>
4217	static LogicalResult verifyPermutationMap(AffineMap permutationMap,
4218	EmitFun emitOpError) {
4219	SmallVector<bool, 8> seen(permutationMap.getNumInputs(), false);
4220	for (auto expr : permutationMap.getResults()) {
4221	auto dim = dyn_cast<AffineDimExpr>(Val&: expr);
4222	auto zero = dyn_cast<AffineConstantExpr>(Val&: expr);
4223	if (zero) {
4224	if (zero.getValue() != 0) {
4225	return emitOpError(
4226	"requires a projected permutation_map (at most one dim or the zero "
4227	"constant can appear in each result)");
4228	}
4229	continue;
4230	}
4231	if (!dim) {
4232	return emitOpError("requires a projected permutation_map (at most one "
4233	"dim or the zero constant can appear in each result)");
4234	}
4235	if (seen[dim.getPosition()]) {
4236	return emitOpError(
4237	"requires a permutation_map that is a permutation (found one dim "
4238	"used more than once)");
4239	}
4240	seen[dim.getPosition()] = true;
4241	}
4242	return success();
4243	}
4244
4245	static LogicalResult
4246	verifyTransferOp(VectorTransferOpInterface op, ShapedType shapedType,
4247	VectorType vectorType, VectorType maskType,
4248	VectorType inferredMaskType, AffineMap permutationMap,
4249	ArrayAttr inBounds) {
4250	if (op->hasAttr("masked")) {
4251	return op->emitOpError("masked attribute has been removed. "
4252	"Use in_bounds instead.");
4253	}
4254
4255	if (!llvm::isa<MemRefType, RankedTensorType>(shapedType))
4256	return op->emitOpError(
4257	"requires source to be a memref or ranked tensor type");
4258
4259	auto elementType = shapedType.getElementType();
4260	DataLayout dataLayout = DataLayout::closest(op: op);
4261	if (auto vectorElementType = llvm::dyn_cast<VectorType>(elementType)) {
4262	// Memref or tensor has vector element type.
4263	unsigned sourceVecSize =
4264	dataLayout.getTypeSizeInBits(t: vectorElementType.getElementType()) *
4265	vectorElementType.getShape().back();
4266	unsigned resultVecSize =
4267	dataLayout.getTypeSizeInBits(t: vectorType.getElementType()) *
4268	vectorType.getShape().back();
4269	if (resultVecSize % sourceVecSize != 0)
4270	return op->emitOpError(
4271	"requires the bitwidth of the minor 1-D vector to be an integral "
4272	"multiple of the bitwidth of the minor 1-D vector of the source");
4273
4274	unsigned sourceVecEltRank = vectorElementType.getRank();
4275	unsigned resultVecRank = vectorType.getRank();
4276	if (sourceVecEltRank > resultVecRank)
4277	return op->emitOpError(
4278	"requires source vector element and vector result ranks to match.");
4279	unsigned rankOffset = resultVecRank - sourceVecEltRank;
4280	// Check that permutation map results match 'rankOffset' of vector type.
4281	if (permutationMap.getNumResults() != rankOffset)
4282	return op->emitOpError("requires a permutation_map with result dims of "
4283	"the same rank as the vector type");
4284
4285	if (maskType)
4286	return op->emitOpError("does not support masks with vector element type");
4287	} else {
4288	// Memref or tensor has scalar element type.
4289	unsigned minorSize =
4290	vectorType.getRank() == 0 ? 1 : vectorType.getShape().back();
4291	unsigned resultVecSize =
4292	dataLayout.getTypeSizeInBits(t: vectorType.getElementType()) * minorSize;
4293	if (resultVecSize % dataLayout.getTypeSizeInBits(t: elementType) != 0)
4294	return op->emitOpError(
4295	"requires the bitwidth of the minor 1-D vector to be an integral "
4296	"multiple of the bitwidth of the source element type");
4297
4298	// Check that permutation map results match rank of vector type.
4299	if (permutationMap.getNumResults() != vectorType.getRank())
4300	return op->emitOpError("requires a permutation_map with result dims of "
4301	"the same rank as the vector type");
4302	}
4303
4304	if (permutationMap.getNumSymbols() != 0)
4305	return op->emitOpError("requires permutation_map without symbols");
4306
4307	if (permutationMap.getNumInputs() != shapedType.getRank())
4308	return op->emitOpError("requires a permutation_map with input dims of the "
4309	"same rank as the source type");
4310
4311	if (maskType && maskType != inferredMaskType)
4312	return op->emitOpError("inferred mask type (")
4313	<< inferredMaskType << ") and mask operand type (" << maskType
4314	<< ") don't match";
4315
4316	if (permutationMap.getNumResults() != static_cast<int64_t>(inBounds.size()))
4317	return op->emitOpError("expects the in_bounds attr of same rank "
4318	"as permutation_map results: ")
4319	<< AffineMapAttr::get(permutationMap)
4320	<< " vs inBounds of size: " << inBounds.size();
4321
4322	return success();
4323	}
4324
4325	static void printTransferAttrs(OpAsmPrinter &p, VectorTransferOpInterface op) {
4326	SmallVector<StringRef, 3> elidedAttrs;
4327	elidedAttrs.push_back(TransferReadOp::getOperandSegmentSizeAttr());
4328	if (op.getPermutationMap().isMinorIdentity())
4329	elidedAttrs.push_back(Elt: op.getPermutationMapAttrName());
4330	// Elide in_bounds attribute if all dims are out-of-bounds.
4331	if (llvm::none_of(op.getInBoundsValues(), [](bool b) { return b; }))
4332	elidedAttrs.push_back(Elt: op.getInBoundsAttrName());
4333	p.printOptionalAttrDict(attrs: op->getAttrs(), elidedAttrs);
4334	}
4335
4336	void TransferReadOp::print(OpAsmPrinter &p) {
4337	p << " " << getBase() << "[" << getIndices() << "], " << getPadding();
4338	if (getMask())
4339	p << ", " << getMask();
4340	printTransferAttrs(p, *this);
4341	p << " : " << getShapedType() << ", " << getVectorType();
4342	}
4343
4344	VectorType mlir::vector::inferTransferOpMaskType(VectorType vecType,
4345	AffineMap permMap) {
4346	auto i1Type = IntegerType::get(permMap.getContext(), 1);
4347	AffineMap invPermMap = inversePermutation(map: compressUnusedDims(map: permMap));
4348	assert(invPermMap && "Inversed permutation map couldn't be computed");
4349	SmallVector<int64_t, 8> maskShape = invPermMap.compose(vecType.getShape());
4350
4351	// The MaskOp specification doesn't support 0-D vectors at the moment. Turn a
4352	// 0-D mask into a single-element 1-D mask.
4353	if (maskShape.empty())
4354	maskShape.push_back(Elt: 1);
4355
4356	SmallVector<bool> scalableDims =
4357	applyPermutationMap(invPermMap, vecType.getScalableDims());
4358
4359	return VectorType::get(maskShape, i1Type, scalableDims);
4360	}
4361
4362	ParseResult TransferReadOp::parse(OpAsmParser &parser, OperationState &result) {
4363	auto &builder = parser.getBuilder();
4364	SMLoc typesLoc;
4365	OpAsmParser::UnresolvedOperand sourceInfo;
4366	SmallVector<OpAsmParser::UnresolvedOperand, 8> indexInfo;
4367	OpAsmParser::UnresolvedOperand paddingInfo;
4368	SmallVector<Type, 2> types;
4369	OpAsmParser::UnresolvedOperand maskInfo;
4370	// Parsing with support for paddingValue.
4371	if (parser.parseOperand(sourceInfo) ||
4372	parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square) ||
4373	parser.parseComma() || parser.parseOperand(paddingInfo))
4374	return failure();
4375	ParseResult hasMask = parser.parseOptionalComma();
4376	if (hasMask.succeeded()) {
4377	if (parser.parseOperand(maskInfo))
4378	return failure();
4379	}
4380	if (parser.parseOptionalAttrDict(result.attributes) ||
4381	parser.getCurrentLocation(&typesLoc) || parser.parseColonTypeList(types))
4382	return failure();
4383	if (types.size() != 2)
4384	return parser.emitError(typesLoc, "requires two types");
4385	auto indexType = builder.getIndexType();
4386	auto shapedType = llvm::dyn_cast<ShapedType>(types[0]);
4387	if (!shapedType || !llvm::isa<MemRefType, RankedTensorType>(shapedType))
4388	return parser.emitError(typesLoc, "requires memref or ranked tensor type");
4389	VectorType vectorType = llvm::dyn_cast<VectorType>(types[1]);
4390	if (!vectorType)
4391	return parser.emitError(typesLoc, "requires vector type");
4392	auto permMapAttrName = TransferReadOp::getPermutationMapAttrName(result.name);
4393	Attribute permMapAttr = result.attributes.get(permMapAttrName);
4394	AffineMap permMap;
4395	if (!permMapAttr) {
4396	if (shapedType.getRank() <
4397	getEffectiveVectorRankForXferOp(shapedType, vectorType))
4398	return parser.emitError(typesLoc,
4399	"expected a custom permutation_map when "
4400	"rank(source) != rank(destination)");
4401	permMap = getTransferMinorIdentityMap(shapedType, vectorType);
4402	result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
4403	} else {
4404	permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
4405	}
4406	auto inBoundsAttrName = TransferReadOp::getInBoundsAttrName(result.name);
4407	Attribute inBoundsAttr = result.attributes.get(inBoundsAttrName);
4408	if (!inBoundsAttr) {
4409	result.addAttribute(inBoundsAttrName,
4410	builder.getBoolArrayAttr(
4411	SmallVector<bool>(permMap.getNumResults(), false)));
4412	}
4413	if (parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
4414	parser.resolveOperands(indexInfo, indexType, result.operands) ||
4415	parser.resolveOperand(paddingInfo, shapedType.getElementType(),
4416	result.operands))
4417	return failure();
4418	if (hasMask.succeeded()) {
4419	if (llvm::dyn_cast<VectorType>(shapedType.getElementType()))
4420	return parser.emitError(
4421	maskInfo.location, "does not support masks with vector element type");
4422	if (vectorType.getRank() != permMap.getNumResults()) {
4423	return parser.emitError(typesLoc,
4424	"expected the same rank for the vector and the "
4425	"results of the permutation map");
4426	}
4427	// Instead of adding the mask type as an op type, compute it based on the
4428	// vector type and the permutation map (to keep the type signature small).
4429	auto maskType = inferTransferOpMaskType(vectorType, permMap);
4430	if (parser.resolveOperand(maskInfo, maskType, result.operands))
4431	return failure();
4432	}
4433	result.addAttribute(TransferReadOp::getOperandSegmentSizeAttr(),
4434	builder.getDenseI32ArrayAttr(
4435	{1, static_cast<int32_t>(indexInfo.size()), 1,
4436	static_cast<int32_t>(hasMask.succeeded())}));
4437	return parser.addTypeToList(vectorType, result.types);
4438	}
4439
4440	LogicalResult TransferReadOp::verify() {
4441	// Consistency of elemental types in source and vector.
4442	ShapedType shapedType = getShapedType();
4443	VectorType vectorType = getVectorType();
4444	VectorType maskType = getMaskType();
4445	auto paddingType = getPadding().getType();
4446	auto permutationMap = getPermutationMap();
4447	VectorType inferredMaskType =
4448	maskType ? inferTransferOpMaskType(vectorType, permutationMap)
4449	: VectorType();
4450	auto sourceElementType = shapedType.getElementType();
4451
4452	if (static_cast<int64_t>(getIndices().size()) != shapedType.getRank())
4453	return emitOpError("requires ") << shapedType.getRank() << " indices";
4454
4455	if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
4456	shapedType, vectorType, maskType,
4457	inferredMaskType, permutationMap, getInBounds())))
4458	return failure();
4459
4460	if (auto sourceVectorElementType =
4461	llvm::dyn_cast<VectorType>(sourceElementType)) {
4462	// Source has vector element type.
4463	// Check that 'sourceVectorElementType' and 'paddingType' types match.
4464	if (sourceVectorElementType != paddingType)
4465	return emitOpError(
4466	"requires source element type and padding type to match.");
4467
4468	} else {
4469	// Check that 'paddingType' is valid to store in a vector type.
4470	if (!VectorType::isValidElementType(paddingType))
4471	return emitOpError("requires valid padding vector elemental type");
4472
4473	// Check that padding type and vector element types match.
4474	if (paddingType != sourceElementType)
4475	return emitOpError(
4476	"requires formal padding and source of the same elemental type");
4477	}
4478
4479	return verifyPermutationMap(permutationMap,
4480	[&](Twine t) { return emitOpError(t); });
4481	}
4482
4483	// MaskableOpInterface methods.
4484
4485	/// Returns the mask type expected by this operation. Mostly used for
4486	/// verification purposes. It requires the operation to be vectorized."
4487	Type TransferReadOp::getExpectedMaskType() {
4488	return inferTransferOpMaskType(getVectorType(), getPermutationMap());
4489	}
4490
4491	//===----------------------------------------------------------------------===//
4492	// TransferReadOp: VectorTransferOpInterface methods.
4493	//===----------------------------------------------------------------------===//
4494	VectorType TransferReadOp::getVectorType() {
4495	return cast<VectorType>(getVector().getType());
4496	}
4497
4498	template <typename TransferOp>
4499	static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
4500	// TODO: support more aggressive createOrFold on:
4501	// op.getIndices()[indicesIdx] + vectorType < dim(op.getSource(), indicesIdx)
4502	if (op.getShapedType().isDynamicDim(indicesIdx))
4503	return false;
4504	Value index = op.getIndices()[indicesIdx];
4505	std::optional<int64_t> cstOp = getConstantIntValue(ofr: index);
4506	if (!cstOp.has_value())
4507	return false;
4508
4509	int64_t sourceSize = op.getShapedType().getDimSize(indicesIdx);
4510	int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
4511
4512	return cstOp.value() + vectorSize <= sourceSize;
4513	}
4514
4515	template <typename TransferOp>
4516	static LogicalResult foldTransferInBoundsAttribute(TransferOp op) {
4517	// TODO: support 0-d corner case.
4518	// TODO: Be less conservative.
4519	if (op.getTransferRank() == 0)
4520	return failure();
4521	AffineMap permutationMap = op.getPermutationMap();
4522	bool changed = false;
4523	SmallVector<bool, 4> newInBounds;
4524	newInBounds.reserve(N: op.getTransferRank());
4525	// Idxs of non-bcast dims - used when analysing bcast dims.
4526	SmallVector<unsigned> nonBcastDims;
4527
4528	// 1. Process non-broadcast dims
4529	for (unsigned i = 0; i < op.getTransferRank(); ++i) {
4530	// 1.1. Already marked as in-bounds, nothing to see here.
4531	if (op.isDimInBounds(i)) {
4532	newInBounds.push_back(Elt: true);
4533	continue;
4534	}
4535	// 1.2. Currently out-of-bounds, check whether we can statically determine
4536	// it is inBounds.
4537	bool inBounds = false;
4538	auto dimExpr = dyn_cast<AffineDimExpr>(Val: permutationMap.getResult(idx: i));
4539	if (dimExpr) {
4540	inBounds = isInBounds(op, /*resultIdx=*/i,
4541	/*indicesIdx=*/dimExpr.getPosition());
4542	nonBcastDims.push_back(Elt: i);
4543	}
4544
4545	newInBounds.push_back(Elt: inBounds);
4546	// We commit the pattern if it is "more inbounds".
4547	changed |= inBounds;
4548	}
4549
4550	// 2. Handle broadcast dims
4551	// If all non-broadcast dims are "in bounds", then all bcast dims should be
4552	// "in bounds" as well.
4553	bool allNonBcastDimsInBounds = llvm::all_of(
4554	nonBcastDims, [&newInBounds](unsigned idx) { return newInBounds[idx]; });
4555	if (allNonBcastDimsInBounds) {
4556	for (size_t idx : permutationMap.getBroadcastDims()) {
4557	changed |= !newInBounds[idx];
4558	newInBounds[idx] = true;
4559	}
4560	}
4561
4562	if (!changed)
4563	return failure();
4564	// OpBuilder is only used as a helper to build an I64ArrayAttr.
4565	OpBuilder b(op.getContext());
4566	op.setInBoundsAttr(b.getBoolArrayAttr(newInBounds));
4567	return success();
4568	}
4569
4570	template <typename TransferOp>
4571	static LogicalResult foldTransferFullMask(TransferOp op) {
4572	auto mask = op.getMask();
4573	if (!mask)
4574	return failure();
4575
4576	if (getMaskFormat(mask) != MaskFormat::AllTrue)
4577	return failure();
4578
4579	op.getMaskMutable().clear();
4580	return success();
4581	}
4582
4583	/// ```
4584	/// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
4585	/// : vector<1x4xf32>, tensor<4x4xf32>
4586	/// %0 = vector.transfer_read %w0[%c1, %c0], %cf0 {in_bounds = [true, true]}
4587	/// : tensor<4x4xf32>, vector<1x4xf32>
4588	/// ```
4589	/// -> Folds into
4590	/// ```
4591	/// %v0
4592	/// ```
4593	static Value foldRAW(TransferReadOp readOp) {
4594	if (!llvm::isa<RankedTensorType>(readOp.getShapedType()))
4595	return {};
4596	auto defWrite = readOp.getBase().getDefiningOp<vector::TransferWriteOp>();
4597	while (defWrite) {
4598	if (checkSameValueRAW(defWrite, readOp))
4599	return defWrite.getVector();
4600	if (!isDisjointTransferIndices(
4601	cast<VectorTransferOpInterface>(defWrite.getOperation()),
4602	cast<VectorTransferOpInterface>(readOp.getOperation())))
4603	break;
4604	defWrite = defWrite.getBase().getDefiningOp<vector::TransferWriteOp>();
4605	}
4606	return {};
4607	}
4608
4609	OpFoldResult TransferReadOp::fold(FoldAdaptor) {
4610	if (Value vec = foldRAW(*this))
4611	return vec;
4612	/// transfer_read(memrefcast) -> transfer_read
4613	if (succeeded(foldTransferInBoundsAttribute(*this)))
4614	return getResult();
4615	if (succeeded(foldTransferFullMask(*this)))
4616	return getResult();
4617	if (succeeded(memref::foldMemRefCast(*this)))
4618	return getResult();
4619	if (succeeded(tensor::foldTensorCast(*this)))
4620	return getResult();
4621	return OpFoldResult();
4622	}
4623
4624	std::optional<SmallVector<int64_t, 4>> TransferReadOp::getShapeForUnroll() {
4625	return llvm::to_vector<4>(getVectorType().getShape());
4626	}
4627
4628	void TransferReadOp::getEffects(
4629	SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
4630	&effects) {
4631	if (llvm::isa<MemRefType>(getShapedType()))
4632	effects.emplace_back(MemoryEffects::Read::get(), &getBaseMutable(),
4633	SideEffects::DefaultResource::get());
4634	}
4635
4636	Speculation::Speculatability TransferReadOp::getSpeculatability() {
4637	if (hasPureTensorSemantics())
4638	return Speculation::Speculatable;
4639	return Speculation::NotSpeculatable;
4640	}
4641
4642	namespace {
4643	/// Store to load forwarding for transfer operations with permuation maps.
4644	/// Even if the permutation maps are different we can still propagate the store
4645	/// into the load if the size of the dimensions read and written match. Then we
4646	/// can replace the transfer_read + transfer_write by vector.broadcast and
4647	/// vector.transpose.
4648	/// Example:
4649	/// ```
4650	/// %w0 = vector.transfer_write %v0, %arg0[%c0, %c0, %c0]
4651	/// {in_bounds = [true, true],
4652	/// permutation_map = affine_map<(d0, d1, d2) -> (d2, d1)>} :
4653	/// vector<4x1xf32>, tensor<4x4x4xf32>
4654	/// %r = vector.transfer_read %w0[%c0, %c0, %c0], %cf0
4655	/// {in_bounds = [true, true, true, true],
4656	/// permutation_map = affine_map<(d0, d1, d2) -> (d1, 0, d2, 0)>} :
4657	/// tensor<4x4x4xf32>, vector<1x100x4x5xf32>
4658	/// ```
4659	/// To:
4660	/// ```
4661	/// %0 = vector.broadcast %arg1 : vector<4x1xf32> to vector<100x5x4x1xf32>
4662	/// %r = vector.transpose %0, [3, 0, 2, 1] :
4663	/// vector<100x5x4x1xf32> to vector<1x100x4x5xf32>
4664	/// ```
4665	struct TransferReadAfterWriteToBroadcast
4666	: public OpRewritePattern<TransferReadOp> {
4667	using OpRewritePattern::OpRewritePattern;
4668
4669	LogicalResult matchAndRewrite(TransferReadOp readOp,
4670	PatternRewriter &rewriter) const override {
4671	if (readOp.hasOutOfBoundsDim() ||
4672	!llvm::isa<RankedTensorType>(readOp.getShapedType()))
4673	return failure();
4674	auto defWrite = readOp.getBase().getDefiningOp<vector::TransferWriteOp>();
4675	if (!defWrite)
4676	return failure();
4677	// TODO: If the written transfer chunk is a superset of the read transfer
4678	// chunk we could do an extract_strided_slice.
4679	if (readOp.getTransferChunkAccessed() !=
4680	defWrite.getTransferChunkAccessed())
4681	return failure();
4682	// TODO: Support cases where a dim is explicitly written but implicitly
4683	// read (i.e., a unit dim that is rank reduced).
4684	if (getUnusedDimsBitVector({readOp.getPermutationMap()}) !=
4685	getUnusedDimsBitVector({defWrite.getPermutationMap()}))
4686	return failure();
4687	if (readOp.getIndices() != defWrite.getIndices() ||
4688	readOp.getMask() != defWrite.getMask())
4689	return failure();
4690	Value vec = defWrite.getVector();
4691	// TODO: loop through the chain of transfer_write if we can prove that they
4692	// don't overlap with the transfer_read. This requires improving
4693	// `isDisjointTransferIndices` helper.
4694	AffineMap readMap = compressUnusedDims(readOp.getPermutationMap());
4695	AffineMap writeMap = compressUnusedDims(defWrite.getPermutationMap());
4696	AffineMap map = readMap.compose(map: writeMap);
4697	if (map.getNumResults() == 0)
4698	return failure();
4699	// Calculate the permutation to apply to go from the vector stored to the
4700	// vector read.
4701	SmallVector<unsigned> permutation;
4702	if (!map.isPermutationOfMinorIdentityWithBroadcasting(permutedDims&: permutation))
4703	return failure();
4704
4705	Location loc = readOp.getLoc();
4706	// Calculate the broadcast shape by applying the reverse permutation to the
4707	// final shape we want.
4708	ArrayRef<int64_t> destShape = readOp.getVectorType().getShape();
4709	SmallVector<int64_t> broadcastShape(destShape.size());
4710	SmallVector<bool> broadcastScalableFlags(destShape.size());
4711	for (const auto &pos : llvm::enumerate(First&: permutation)) {
4712	broadcastShape[pos.value()] = destShape[pos.index()];
4713	broadcastScalableFlags[pos.value()] =
4714	readOp.getVectorType().getScalableDims()[pos.index()];
4715	}
4716	VectorType broadcastedType = VectorType::get(
4717	broadcastShape, defWrite.getVectorType().getElementType(),
4718	broadcastScalableFlags);
4719	vec = rewriter.create<vector::BroadcastOp>(loc, broadcastedType, vec);
4720	SmallVector<int64_t> transposePerm(permutation.begin(), permutation.end());
4721	rewriter.replaceOpWithNewOp<vector::TransposeOp>(readOp, vec,
4722	transposePerm);
4723	return success();
4724	}
4725	};
4726	} // namespace
4727
4728	void TransferReadOp::getCanonicalizationPatterns(RewritePatternSet &results,
4729	MLIRContext *context) {
4730	results.add<TransferReadAfterWriteToBroadcast>(context);
4731	}
4732
4733	//===----------------------------------------------------------------------===//
4734	// TransferWriteOp
4735	//===----------------------------------------------------------------------===//
4736
4737	/// 1. Builder with type inference.
4738	void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4739	Value vector, Value dest, ValueRange indices,
4740	AffineMapAttr permutationMapAttr,
4741	/*optional*/ Value mask,
4742	/*optional*/ ArrayAttr inBoundsAttr) {
4743	Type resultType = llvm::dyn_cast<RankedTensorType>(dest.getType());
4744	build(builder, result, resultType, vector, dest, indices, permutationMapAttr,
4745	mask, inBoundsAttr);
4746	}
4747
4748	/// 2. Builder with type inference that sets an empty mask (variant with attrs).
4749	void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4750	Value vector, Value dest, ValueRange indices,
4751	AffineMapAttr permutationMapAttr,
4752	/*optional*/ ArrayAttr inBoundsAttr) {
4753	build(builder, result, vector, dest, indices, permutationMapAttr,
4754	/*mask=*/Value(), inBoundsAttr);
4755	}
4756
4757	/// 3. Builder with type inference that sets an empty mask (variant without
4758	/// attrs)
4759	void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4760	Value vector, Value dest, ValueRange indices,
4761	AffineMap permutationMap,
4762	std::optional<ArrayRef<bool>> inBounds) {
4763	auto permutationMapAttr = AffineMapAttr::get(permutationMap);
4764	auto inBoundsAttr =
4765	(inBounds && !inBounds.value().empty())
4766	? builder.getBoolArrayAttr(inBounds.value())
4767	: builder.getBoolArrayAttr(SmallVector<bool>(
4768	llvm::cast<VectorType>(vector.getType()).getRank(), false));
4769	build(builder, result, vector, dest, indices, permutationMapAttr,
4770	/*mask=*/Value(), inBoundsAttr);
4771	}
4772
4773	/// 4. Builder with type inference that sets an empty mask and sets permutation
4774	/// map to 'getMinorIdentityMap'.
4775	void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
4776	Value vector, Value dest, ValueRange indices,
4777	std::optional<ArrayRef<bool>> inBounds) {
4778	auto vectorType = llvm::cast<VectorType>(vector.getType());
4779	AffineMap permutationMap = getTransferMinorIdentityMap(
4780	llvm::cast<ShapedType>(dest.getType()), vectorType);
4781	build(builder, result, vector, dest, indices, permutationMap, inBounds);
4782	}
4783
4784	ParseResult TransferWriteOp::parse(OpAsmParser &parser,
4785	OperationState &result) {
4786	auto &builder = parser.getBuilder();
4787	SMLoc typesLoc;
4788	OpAsmParser::UnresolvedOperand vectorInfo, sourceInfo;
4789	SmallVector<OpAsmParser::UnresolvedOperand, 8> indexInfo;
4790	SmallVector<Type, 2> types;
4791	OpAsmParser::UnresolvedOperand maskInfo;
4792	if (parser.parseOperand(vectorInfo) || parser.parseComma() ||
4793	parser.parseOperand(sourceInfo) ||
4794	parser.parseOperandList(indexInfo, OpAsmParser::Delimiter::Square))
4795	return failure();
4796	ParseResult hasMask = parser.parseOptionalComma();
4797	if (hasMask.succeeded() && parser.parseOperand(maskInfo))
4798	return failure();
4799	if (parser.parseOptionalAttrDict(result.attributes) ||
4800	parser.getCurrentLocation(&typesLoc) || parser.parseColonTypeList(types))
4801	return failure();
4802	if (types.size() != 2)
4803	return parser.emitError(typesLoc, "requires two types");
4804	auto indexType = builder.getIndexType();
4805	VectorType vectorType = llvm::dyn_cast<VectorType>(types[0]);
4806	if (!vectorType)
4807	return parser.emitError(typesLoc, "requires vector type");
4808	ShapedType shapedType = llvm::dyn_cast<ShapedType>(types[1]);
4809	if (!shapedType || !llvm::isa<MemRefType, RankedTensorType>(shapedType))
4810	return parser.emitError(typesLoc, "requires memref or ranked tensor type");
4811	auto permMapAttrName =
4812	TransferWriteOp::getPermutationMapAttrName(result.name);
4813	auto permMapAttr = result.attributes.get(permMapAttrName);
4814	AffineMap permMap;
4815	if (!permMapAttr) {
4816	if (shapedType.getRank() <
4817	getEffectiveVectorRankForXferOp(shapedType, vectorType))
4818	return parser.emitError(typesLoc,
4819	"expected a custom permutation_map when "
4820	"rank(source) != rank(destination)");
4821	permMap = getTransferMinorIdentityMap(shapedType, vectorType);
4822	result.attributes.set(permMapAttrName, AffineMapAttr::get(permMap));
4823	} else {
4824	permMap = llvm::cast<AffineMapAttr>(permMapAttr).getValue();
4825	}
4826	auto inBoundsAttrName = TransferWriteOp::getInBoundsAttrName(result.name);
4827	Attribute inBoundsAttr = result.attributes.get(inBoundsAttrName);
4828	if (!inBoundsAttr) {
4829	result.addAttribute(inBoundsAttrName,
4830	builder.getBoolArrayAttr(
4831	SmallVector<bool>(permMap.getNumResults(), false)));
4832	}
4833	if (parser.resolveOperand(vectorInfo, vectorType, result.operands) ||
4834	parser.resolveOperand(sourceInfo, shapedType, result.operands) ||
4835	parser.resolveOperands(indexInfo, indexType, result.operands))
4836	return failure();
4837	if (hasMask.succeeded()) {
4838	if (llvm::dyn_cast<VectorType>(shapedType.getElementType()))
4839	return parser.emitError(
4840	maskInfo.location, "does not support masks with vector element type");
4841	if (vectorType.getRank() != permMap.getNumResults()) {
4842	return parser.emitError(typesLoc,
4843	"expected the same rank for the vector and the "
4844	"results of the permutation map");
4845	}
4846	auto maskType = inferTransferOpMaskType(vectorType, permMap);
4847	if (parser.resolveOperand(maskInfo, maskType, result.operands))
4848	return failure();
4849	}
4850	result.addAttribute(TransferWriteOp::getOperandSegmentSizeAttr(),
4851	builder.getDenseI32ArrayAttr(
4852	{1, 1, static_cast<int32_t>(indexInfo.size()),
4853	static_cast<int32_t>(hasMask.succeeded())}));
4854	return failure(llvm::isa<RankedTensorType>(shapedType) &&
4855	parser.addTypeToList(shapedType, result.types));
4856	}
4857
4858	void TransferWriteOp::print(OpAsmPrinter &p) {
4859	p << " " << getVector() << ", " << getBase() << "[" << getIndices() << "]";
4860	if (getMask())
4861	p << ", " << getMask();
4862	printTransferAttrs(p, *this);
4863	p << " : " << getVectorType() << ", " << getShapedType();
4864	}
4865
4866	LogicalResult TransferWriteOp::verify() {
4867	// Consistency of elemental types in shape and vector.
4868	ShapedType shapedType = getShapedType();
4869	VectorType vectorType = getVectorType();
4870	VectorType maskType = getMaskType();
4871	auto permutationMap = getPermutationMap();
4872	VectorType inferredMaskType =
4873	maskType ? inferTransferOpMaskType(vectorType, permutationMap)
4874	: VectorType();
4875
4876	if (llvm::size(getIndices()) != shapedType.getRank())
4877	return emitOpError("requires ") << shapedType.getRank() << " indices";
4878
4879	// We do not allow broadcast dimensions on TransferWriteOps for the moment,
4880	// as the semantics is unclear. This can be revisited later if necessary.
4881	if (hasBroadcastDim())
4882	return emitOpError("should not have broadcast dimensions");
4883
4884	if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
4885	shapedType, vectorType, maskType,
4886	inferredMaskType, permutationMap, getInBounds())))
4887	return failure();
4888
4889	return verifyPermutationMap(permutationMap,
4890	[&](Twine t) { return emitOpError(t); });
4891	}
4892
4893	//===----------------------------------------------------------------------===//
4894	// TransferWriteOp: MaskableOpInterface methods.
4895	//===----------------------------------------------------------------------===//
4896
4897	/// Returns the mask type expected by this operation. Mostly used for
4898	/// verification purposes.
4899	Type TransferWriteOp::getExpectedMaskType() {
4900	return inferTransferOpMaskType(getVectorType(), getPermutationMap());
4901	}
4902
4903	//===----------------------------------------------------------------------===//
4904	// TransferWriteOp: VectorTransferOpInterface methods.
4905	//===----------------------------------------------------------------------===//
4906	Value TransferWriteOp::getVector() { return getOperand(0); }
4907	VectorType TransferWriteOp::getVectorType() {
4908	return cast<VectorType>(getValueToStore().getType());
4909	}
4910
4911	//===----------------------------------------------------------------------===//
4912	// TransferWriteOp: fold methods.
4913	//===----------------------------------------------------------------------===//
4914	/// Fold:
4915	/// ```
4916	/// %t1 = ...
4917	/// %v = vector.transfer_read %t0[%c0...], {in_bounds = [true...]} :
4918	/// tensor<static_sizesxf32>, vector<static_sizesxf32>
4919	/// %t2 = vector.transfer_write %v, %t1[%c0...] {in_bounds = [true...]} :
4920	/// vector<static_sizesxf32>, tensor<static_sizesxf32>
4921	/// ```
4922	///
4923	/// into:
4924	///
4925	/// ```
4926	/// %t0
4927	/// ```
4928	///
4929	/// The producer of t1 may or may not be DCE'd depending on whether it is a
4930	/// block argument or has side effects.
4931	static LogicalResult foldReadInitWrite(TransferWriteOp write,
4932	ArrayRef<Attribute>,
4933	SmallVectorImpl<OpFoldResult> &results) {
4934	// TODO: support 0-d corner case.
4935	if (write.getTransferRank() == 0)
4936	return failure();
4937	auto rankedTensorType =
4938	llvm::dyn_cast<RankedTensorType>(write.getBase().getType());
4939	// If not operating on tensors, bail.
4940	if (!rankedTensorType)
4941	return failure();
4942	// If no read, bail.
4943	auto read = write.getVector().getDefiningOp<vector::TransferReadOp>();
4944	if (!read)
4945	return failure();
4946	// TODO: support 0-d corner case.
4947	if (read.getTransferRank() == 0)
4948	return failure();
4949	// For now, only accept minor identity. Future: composition is minor identity.
4950	if (!read.getPermutationMap().isMinorIdentity() ||
4951	!write.getPermutationMap().isMinorIdentity())
4952	return failure();
4953	// Bail on mismatching ranks.
4954	if (read.getTransferRank() != write.getTransferRank())
4955	return failure();
4956	// Bail on potential out-of-bounds accesses.
4957	if (read.hasOutOfBoundsDim() || write.hasOutOfBoundsDim())
4958	return failure();
4959	// Tensor types must be the same.
4960	if (read.getBase().getType() != rankedTensorType)
4961	return failure();
4962	// Vector types must be the same.
4963	if (read.getVectorType() != write.getVectorType())
4964	return failure();
4965	// Vector and Tensor shapes must match.
4966	if (read.getVectorType().getShape() != rankedTensorType.getShape())
4967	return failure();
4968	// If any index is nonzero.
4969	auto isNotConstantZero = [](Value v) {
4970	auto cstOp = getConstantIntValue(ofr: v);
4971	return !cstOp.has_value() || cstOp.value() != 0;
4972	};
4973	if (llvm::any_of(read.getIndices(), isNotConstantZero) ||
4974	llvm::any_of(write.getIndices(), isNotConstantZero))
4975	return failure();
4976	// Success.
4977	results.push_back(Elt: read.getBase());
4978	return success();
4979	}
4980
4981	static bool checkSameValueWAR(vector::TransferReadOp read,
4982	vector::TransferWriteOp write) {
4983	return read.getBase() == write.getBase() &&
4984	read.getIndices() == write.getIndices() &&
4985	read.getPermutationMap() == write.getPermutationMap() &&
4986	read.getVectorType() == write.getVectorType() && !read.getMask() &&
4987	!write.getMask();
4988	}
4989	/// Fold transfer_write write after read:
4990	/// ```
4991	/// %t0 = ...
4992	/// %v = vector.transfer_read %t0[%c0...] :
4993	/// tensor<static_sizesxf32>, vector<static_sizesxf32>
4994	/// %t1 = vector.transfer_write %v, %t0[%c0...] :
4995	/// vector<static_sizesxf32>, tensor<static_sizesxf32>
4996	/// ```
4997	///
4998	/// into:
4999	///
5000	/// ```
5001	/// %t0
5002	/// ```
5003	static LogicalResult foldWAR(TransferWriteOp write,
5004	SmallVectorImpl<OpFoldResult> &results) {
5005	if (!llvm::isa<RankedTensorType>(write.getBase().getType()))
5006	return failure();
5007	auto read = write.getVector().getDefiningOp<vector::TransferReadOp>();
5008	if (!read)
5009	return failure();
5010
5011	if (!checkSameValueWAR(read, write))
5012	return failure();
5013	results.push_back(Elt: read.getBase());
5014	return success();
5015	}
5016
5017	LogicalResult TransferWriteOp::fold(FoldAdaptor adaptor,
5018	SmallVectorImpl<OpFoldResult> &results) {
5019	if (succeeded(foldReadInitWrite(*this, adaptor.getOperands(), results)))
5020	return success();
5021	if (succeeded(foldWAR(*this, results)))
5022	return success();
5023	if (succeeded(foldTransferInBoundsAttribute(*this)))
5024	return success();
5025	if (succeeded(foldTransferFullMask(*this)))
5026	return success();
5027	return memref::foldMemRefCast(*this);
5028	}
5029
5030	//===----------------------------------------------------------------------===//
5031	// TransferWriteOp: other methods.
5032	//===----------------------------------------------------------------------===//
5033	std::optional<SmallVector<int64_t, 4>> TransferWriteOp::getShapeForUnroll() {
5034	return llvm::to_vector<4>(getVectorType().getShape());
5035	}
5036
5037	void TransferWriteOp::getEffects(
5038	SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
5039	&effects) {
5040	if (llvm::isa<MemRefType>(getShapedType()))
5041	effects.emplace_back(MemoryEffects::Write::get(), &getBaseMutable(),
5042	SideEffects::DefaultResource::get());
5043	}
5044
5045	Speculation::Speculatability TransferWriteOp::getSpeculatability() {
5046	if (hasPureTensorSemantics())
5047	return Speculation::Speculatable;
5048	return Speculation::NotSpeculatable;
5049	}
5050
5051	namespace {
5052	/// Remove dead transfer write from the SSA chain so that it an be eliminated by
5053	/// DCE
5054	/// ```
5055	/// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
5056	/// : vector<1x4xf32>, tensor<4x4xf32>
5057	/// %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]}
5058	/// : vector<1x4xf32>, tensor<4x4xf32>
5059	/// %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
5060	/// : vector<1x4xf32>, tensor<4x4xf32>
5061	/// ```
5062	///
5063	/// into:
5064	///
5065	/// ```
5066	/// %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]}
5067	/// : vector<1x4xf32>, tensor<4x4xf32>
5068	/// %w1 = vector.transfer_write %v0, %arg0[%c2, %c0] {in_bounds = [true, true]}
5069	/// : vector<1x4xf32>, tensor<4x4xf32>
5070	/// %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]}
5071	/// : vector<1x4xf32>, tensor<4x4xf32>
5072	/// ```
5073	///
5074	/// `%w0 = vector.transfer_write` op will be removed by DCE if it doesn't have
5075	/// any other uses.
5076	class FoldWaw final : public OpRewritePattern<TransferWriteOp> {
5077	public:
5078	using OpRewritePattern::OpRewritePattern;
5079	LogicalResult matchAndRewrite(TransferWriteOp writeOp,
5080	PatternRewriter &rewriter) const override {
5081	if (!llvm::isa<RankedTensorType>(writeOp.getShapedType()))
5082	return failure();
5083	vector::TransferWriteOp writeToModify = writeOp;
5084
5085	auto defWrite = writeOp.getBase().getDefiningOp<vector::TransferWriteOp>();
5086	while (defWrite) {
5087	if (checkSameValueWAW(writeOp, defWrite)) {
5088	rewriter.modifyOpInPlace(writeToModify, [&]() {
5089	writeToModify.getBaseMutable().assign(defWrite.getBase());
5090	});
5091	return success();
5092	}
5093	if (!isDisjointTransferIndices(
5094	cast<VectorTransferOpInterface>(defWrite.getOperation()),
5095	cast<VectorTransferOpInterface>(writeOp.getOperation())))
5096	break;
5097	// If the previous write op doesn't have any other use we an safely look
5098	// at the previous store to see if it can be removed.
5099	if (!defWrite->hasOneUse())
5100	break;
5101	writeToModify = defWrite;
5102	defWrite = defWrite.getBase().getDefiningOp<vector::TransferWriteOp>();
5103	}
5104	return failure();
5105	}
5106	};
5107
5108	/// Rewrite tensor::ExtractSliceOp(vector::TransferWriteOp) to
5109	/// vector::TransferWriteOp(tensor::ExtractSliceOp) if the full slice is
5110	/// overwritten and inserted into another tensor. After this rewrite, the
5111	/// operations bufferize in-place since all of them work on the same slice.
5112	///
5113	/// For example:
5114	/// ```mlir
5115	/// %0 = vector.transfer_write %vec, %init_tensor[%c0, %c0]
5116	/// : vector<8x16xf32>, tensor<8x16xf32>
5117	/// %1 = tensor.extract_slice %0[0, 0] [%sz0, %sz1] [1, 1]
5118	/// : tensor<8x16xf32> to tensor<?x?xf32>
5119	/// %r = tensor.insert_slice %1 into %iter_arg[%iv0, %iv1] [%sz0, %sz1] [1, 1]
5120	/// : tensor<?x?xf32> into tensor<27x37xf32>
5121	/// ```
5122	/// folds to
5123	/// ```mlir
5124	/// %0 = tensor.extract_slice %iter_arg[%iv0, %iv1] [%sz0, %sz1] [1, 1]
5125	/// : tensor<27x37xf32> to tensor<?x?xf32>
5126	/// %1 = vector.transfer_write %vec, %0[%c0, %c0]
5127	/// : vector<8x16xf32>, tensor<?x?xf32>
5128	/// %r = tensor.insert_slice %1 into %iter_arg[%iv0, %iv1] [%sz0, %sz1] [1, 1]
5129	/// : tensor<?x?xf32> into tensor<27x37xf32>
5130	/// ```
5131	struct SwapExtractSliceOfTransferWrite
5132	: public OpRewritePattern<tensor::InsertSliceOp> {
5133	public:
5134	using OpRewritePattern::OpRewritePattern;
5135
5136	LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
5137	PatternRewriter &rewriter) const override {
5138	if (!insertOp.hasUnitStride())
5139	return failure();
5140	auto extractOp =
5141	insertOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
5142	if (!extractOp || !extractOp.hasUnitStride() || !extractOp->hasOneUse())
5143	return failure();
5144	auto transferOp = extractOp.getSource().getDefiningOp<TransferWriteOp>();
5145	if (!transferOp || !transferOp->hasOneUse())
5146	return failure();
5147
5148	// Fail if vector::TransferWriteOp or tensor::ExtractSliceOp is
5149	// rank-reducing.
5150	if (insertOp.getSourceType().getRank() != transferOp.getTransferRank()) {
5151	return rewriter.notifyMatchFailure(insertOp,
5152	"use-def chain is rank-reducing");
5153	}
5154
5155	// Fail if tensor::ExtractSliceOp has non-zero offset.
5156	if (!extractOp.hasZeroOffset()) {
5157	return rewriter.notifyMatchFailure(insertOp,
5158	"ExtractSliceOp has non-zero offset");
5159	}
5160
5161	// Fail if tensor::TransferWriteOp has non-zero offset.
5162	if (!llvm::all_of(transferOp.getIndices(), [](Value value) {
5163	return getConstantIntValue(ofr: value) == static_cast<int64_t>(0);
5164	})) {
5165	return rewriter.notifyMatchFailure(insertOp,
5166	"TranferWriteOp has non-zero offset");
5167	}
5168
5169	// Fail if tensor::ExtractSliceOp and tensor::InsertSliceOp sizes differ.
5170	if (insertOp.getMixedSizes().size() != extractOp.getMixedSizes().size()) {
5171	return rewriter.notifyMatchFailure(
5172	insertOp, "InsertSliceOp and ExtractSliceOp ranks differ");
5173	}
5174
5175	for (auto [insertSize, extractSize] :
5176	llvm::zip_equal(insertOp.getMixedSizes(), extractOp.getMixedSizes())) {
5177	if (!isEqualConstantIntOrValue(insertSize, extractSize)) {
5178	return rewriter.notifyMatchFailure(
5179	insertOp, "InsertSliceOp and ExtractSliceOp sizes differ");
5180	}
5181	}
5182
5183	// Fail if the vector::TransferWriteOp may not overwrite the full tensor.
5184	assert(transferOp.getVectorType().hasStaticShape() &&
5185	"expected vector to have a static shape");
5186	ArrayRef<int64_t> vectorShape = transferOp.getVectorType().getShape();
5187	SmallVector<int64_t> resultShape = applyPermutationMap(
5188	transferOp.getPermutationMap(), transferOp.getShapedType().getShape());
5189	if (transferOp.getMask() || !vectorShape.equals(RHS: resultShape)) {
5190	return rewriter.notifyMatchFailure(
5191	insertOp, "TransferWriteOp may not write the full tensor.");
5192	}
5193
5194	// Swap the tensor::ExtractSliceOp in front of the vector::TransferWriteOp.
5195	// Set all in_bounds to false and let the folder infer them.
5196	SmallVector<bool> newInBounds(vectorShape.size(), false);
5197	auto newExtractOp = rewriter.create<tensor::ExtractSliceOp>(
5198	extractOp.getLoc(), insertOp.getSourceType(), insertOp.getDest(),
5199	insertOp.getMixedOffsets(), insertOp.getMixedSizes(),
5200	insertOp.getMixedStrides());
5201	auto newTransferWriteOp = rewriter.create<TransferWriteOp>(
5202	transferOp.getLoc(), transferOp.getVector(), newExtractOp.getResult(),
5203	transferOp.getIndices(), transferOp.getPermutationMapAttr(),
5204	rewriter.getBoolArrayAttr(newInBounds));
5205	rewriter.modifyOpInPlace(insertOp, [&]() {
5206	insertOp.getSourceMutable().assign(newTransferWriteOp.getResult());
5207	});
5208	return success();
5209	}
5210	};
5211
5212	} // namespace
5213
5214	void TransferWriteOp::getCanonicalizationPatterns(RewritePatternSet &results,
5215	MLIRContext *context) {
5216	results.add<FoldWaw, SwapExtractSliceOfTransferWrite>(context);
5217	}
5218
5219	//===----------------------------------------------------------------------===//
5220	// LoadOp
5221	//===----------------------------------------------------------------------===//
5222
5223	static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,
5224	VectorType vecTy,
5225	MemRefType memRefTy) {
5226	// If rank==0 or size==1 it's equivalent to scalar load/store, so we don't
5227	// need any strides limitations.
5228	if (!vecTy.isScalable() &&
5229	(vecTy.getRank() == 0 || vecTy.getNumElements() == 1))
5230	return success();
5231
5232	if (!memRefTy.isLastDimUnitStride())
5233	return op->emitOpError(message: "most minor memref dim must have unit stride");
5234	return success();
5235	}
5236
5237	LogicalResult vector::LoadOp::verify() {
5238	VectorType resVecTy = getVectorType();
5239	MemRefType memRefTy = getMemRefType();
5240
5241	if (failed(verifyLoadStoreMemRefLayout(*this, resVecTy, memRefTy)))
5242	return failure();
5243
5244	if (memRefTy.getRank() < resVecTy.getRank())
5245	return emitOpError(
5246	"destination memref has lower rank than the result vector");
5247
5248	// Checks for vector memrefs.
5249	Type memElemTy = memRefTy.getElementType();
5250	if (auto memVecTy = llvm::dyn_cast<VectorType>(memElemTy)) {
5251	if (memVecTy != resVecTy)
5252	return emitOpError("base memref and result vector types should match");
5253	memElemTy = memVecTy.getElementType();
5254	}
5255
5256	if (resVecTy.getElementType() != memElemTy)
5257	return emitOpError("base and result element types should match");
5258	if (llvm::size(getIndices()) != memRefTy.getRank())
5259	return emitOpError("requires ") << memRefTy.getRank() << " indices";
5260	return success();
5261	}
5262
5263	OpFoldResult LoadOp::fold(FoldAdaptor) {
5264	if (succeeded(memref::foldMemRefCast(*this)))
5265	return getResult();
5266	return OpFoldResult();
5267	}
5268
5269	//===----------------------------------------------------------------------===//
5270	// StoreOp
5271	//===----------------------------------------------------------------------===//
5272
5273	LogicalResult vector::StoreOp::verify() {
5274	VectorType valueVecTy = getVectorType();
5275	MemRefType memRefTy = getMemRefType();
5276
5277	if (failed(verifyLoadStoreMemRefLayout(*this, valueVecTy, memRefTy)))
5278	return failure();
5279
5280	if (memRefTy.getRank() < valueVecTy.getRank())
5281	return emitOpError("source memref has lower rank than the vector to store");
5282
5283	// Checks for vector memrefs.
5284	Type memElemTy = memRefTy.getElementType();
5285	if (auto memVecTy = llvm::dyn_cast<VectorType>(memElemTy)) {
5286	if (memVecTy != valueVecTy)
5287	return emitOpError(
5288	"base memref and valueToStore vector types should match");
5289	memElemTy = memVecTy.getElementType();
5290	}
5291
5292	if (valueVecTy.getElementType() != memElemTy)
5293	return emitOpError("base and valueToStore element type should match");
5294	if (llvm::size(getIndices()) != memRefTy.getRank())
5295	return emitOpError("requires ") << memRefTy.getRank() << " indices";
5296	return success();
5297	}
5298
5299	LogicalResult StoreOp::fold(FoldAdaptor adaptor,
5300	SmallVectorImpl<OpFoldResult> &results) {
5301	return memref::foldMemRefCast(*this);
5302	}
5303
5304	//===----------------------------------------------------------------------===//
5305	// MaskedLoadOp
5306	//===----------------------------------------------------------------------===//
5307
5308	LogicalResult MaskedLoadOp::verify() {
5309	VectorType maskVType = getMaskVectorType();
5310	VectorType passVType = getPassThruVectorType();
5311	VectorType resVType = getVectorType();
5312	MemRefType memType = getMemRefType();
5313
5314	if (resVType.getElementType() != memType.getElementType())
5315	return emitOpError("base and result element type should match");
5316	if (llvm::size(getIndices()) != memType.getRank())
5317	return emitOpError("requires ") << memType.getRank() << " indices";
5318	if (resVType.getShape() != maskVType.getShape())
5319	return emitOpError("expected result shape to match mask shape");
5320	if (resVType != passVType)
5321	return emitOpError("expected pass_thru of same type as result type");
5322	return success();
5323	}
5324
5325	namespace {
5326	class MaskedLoadFolder final : public OpRewritePattern<MaskedLoadOp> {
5327	public:
5328	using OpRewritePattern::OpRewritePattern;
5329	LogicalResult matchAndRewrite(MaskedLoadOp load,
5330	PatternRewriter &rewriter) const override {
5331	switch (getMaskFormat(load.getMask())) {
5332	case MaskFormat::AllTrue:
5333	rewriter.replaceOpWithNewOp<vector::LoadOp>(
5334	load, load.getType(), load.getBase(), load.getIndices());
5335	return success();
5336	case MaskFormat::AllFalse:
5337	rewriter.replaceOp(load, load.getPassThru());
5338	return success();
5339	case MaskFormat::Unknown:
5340	return failure();
5341	}
5342	llvm_unreachable("Unexpected 1DMaskFormat on MaskedLoad");
5343	}
5344	};
5345	} // namespace
5346
5347	void MaskedLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
5348	MLIRContext *context) {
5349	results.add<MaskedLoadFolder>(context);
5350	}
5351
5352	OpFoldResult MaskedLoadOp::fold(FoldAdaptor) {
5353	if (succeeded(memref::foldMemRefCast(*this)))
5354	return getResult();
5355	return OpFoldResult();
5356	}
5357
5358	//===----------------------------------------------------------------------===//
5359	// MaskedStoreOp
5360	//===----------------------------------------------------------------------===//
5361
5362	LogicalResult MaskedStoreOp::verify() {
5363	VectorType maskVType = getMaskVectorType();
5364	VectorType valueVType = getVectorType();
5365	MemRefType memType = getMemRefType();
5366
5367	if (valueVType.getElementType() != memType.getElementType())
5368	return emitOpError("base and valueToStore element type should match");
5369	if (llvm::size(getIndices()) != memType.getRank())
5370	return emitOpError("requires ") << memType.getRank() << " indices";
5371	if (valueVType.getShape() != maskVType.getShape())
5372	return emitOpError("expected valueToStore shape to match mask shape");
5373	return success();
5374	}
5375
5376	namespace {
5377	class MaskedStoreFolder final : public OpRewritePattern<MaskedStoreOp> {
5378	public:
5379	using OpRewritePattern::OpRewritePattern;
5380	LogicalResult matchAndRewrite(MaskedStoreOp store,
5381	PatternRewriter &rewriter) const override {
5382	switch (getMaskFormat(store.getMask())) {
5383	case MaskFormat::AllTrue:
5384	rewriter.replaceOpWithNewOp<vector::StoreOp>(
5385	store, store.getValueToStore(), store.getBase(), store.getIndices());
5386	return success();
5387	case MaskFormat::AllFalse:
5388	rewriter.eraseOp(op: store);
5389	return success();
5390	case MaskFormat::Unknown:
5391	return failure();
5392	}
5393	llvm_unreachable("Unexpected 1DMaskFormat on MaskedStore");
5394	}
5395	};
5396	} // namespace
5397
5398	void MaskedStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
5399	MLIRContext *context) {
5400	results.add<MaskedStoreFolder>(context);
5401	}
5402
5403	LogicalResult MaskedStoreOp::fold(FoldAdaptor adaptor,
5404	SmallVectorImpl<OpFoldResult> &results) {
5405	return memref::foldMemRefCast(*this);
5406	}
5407
5408	//===----------------------------------------------------------------------===//
5409	// GatherOp
5410	//===----------------------------------------------------------------------===//
5411
5412	LogicalResult GatherOp::verify() {
5413	VectorType indVType = getIndexVectorType();
5414	VectorType maskVType = getMaskVectorType();
5415	VectorType resVType = getVectorType();
5416	ShapedType baseType = getBaseType();
5417
5418	if (!llvm::isa<MemRefType, RankedTensorType>(baseType))
5419	return emitOpError("requires base to be a memref or ranked tensor type");
5420
5421	if (resVType.getElementType() != baseType.getElementType())
5422	return emitOpError("base and result element type should match");
5423	if (llvm::size(getIndices()) != baseType.getRank())
5424	return emitOpError("requires ") << baseType.getRank() << " indices";
5425	if (resVType.getShape() != indVType.getShape())
5426	return emitOpError("expected result dim to match indices dim");
5427	if (resVType.getShape() != maskVType.getShape())
5428	return emitOpError("expected result dim to match mask dim");
5429	if (resVType != getPassThruVectorType())
5430	return emitOpError("expected pass_thru of same type as result type");
5431	return success();
5432	}
5433
5434	// MaskableOpInterface methods.
5435
5436	/// Returns the mask type expected by this operation. Mostly used for
5437	/// verification purposes. It requires the operation to be vectorized."
5438	Type GatherOp::getExpectedMaskType() {
5439	auto vecType = this->getIndexVectorType();
5440	return VectorType::get(vecType.getShape(),
5441	IntegerType::get(vecType.getContext(), /*width=*/1),
5442	vecType.getScalableDims());
5443	}
5444
5445	std::optional<SmallVector<int64_t, 4>> GatherOp::getShapeForUnroll() {
5446	return llvm::to_vector<4>(getVectorType().getShape());
5447	}
5448
5449	/// Cheeck if `indexVec` is constant 1D vec of consecutive values [0, 1, 2, ...]
5450	static LogicalResult isZeroBasedContiguousSeq(Value indexVec) {
5451	auto vecType = dyn_cast<VectorType>(indexVec.getType());
5452	if (!vecType || vecType.getRank() != 1 || vecType.isScalable())
5453	return failure();
5454
5455	if (indexVec.getDefiningOp<StepOp>())
5456	return success();
5457
5458	DenseIntElementsAttr elements;
5459	if (!matchPattern(value: indexVec, pattern: m_Constant(bind_value: &elements)))
5460	return failure();
5461
5462	return success(
5463	llvm::equal(elements, llvm::seq<int64_t>(0, vecType.getNumElements())));
5464	}
5465
5466	namespace {
5467	class GatherFolder final : public OpRewritePattern<GatherOp> {
5468	public:
5469	using OpRewritePattern::OpRewritePattern;
5470	LogicalResult matchAndRewrite(GatherOp gather,
5471	PatternRewriter &rewriter) const override {
5472	switch (getMaskFormat(gather.getMask())) {
5473	case MaskFormat::AllTrue:
5474	return failure(); // no unmasked equivalent
5475	case MaskFormat::AllFalse:
5476	rewriter.replaceOp(gather, gather.getPassThru());
5477	return success();
5478	case MaskFormat::Unknown:
5479	return failure();
5480	}
5481	llvm_unreachable("Unexpected 1DMaskFormat on GatherFolder");
5482	}
5483	};
5484
5485	/// Fold gathers with consecutive offsets [0, 1, 2, ...] into contiguous
5486	/// maskedload. Only 1D fixed vectors are supported for now.
5487	class FoldContiguousGather final : public OpRewritePattern<GatherOp> {
5488	public:
5489	using OpRewritePattern::OpRewritePattern;
5490	LogicalResult matchAndRewrite(GatherOp op,
5491	PatternRewriter &rewriter) const override {
5492	if (!isa<MemRefType>(op.getBase().getType()))
5493	return rewriter.notifyMatchFailure(op, "base must be of memref type");
5494
5495	if (failed(isZeroBasedContiguousSeq(op.getIndexVec())))
5496	return failure();
5497
5498	rewriter.replaceOpWithNewOp<MaskedLoadOp>(op, op.getType(), op.getBase(),
5499	op.getIndices(), op.getMask(),
5500	op.getPassThru());
5501	return success();
5502	}
5503	};
5504	} // namespace
5505
5506	void GatherOp::getCanonicalizationPatterns(RewritePatternSet &results,
5507	MLIRContext *context) {
5508	results.add<GatherFolder, FoldContiguousGather>(context);
5509	}
5510
5511	//===----------------------------------------------------------------------===//
5512	// ScatterOp
5513	//===----------------------------------------------------------------------===//
5514
5515	LogicalResult ScatterOp::verify() {
5516	VectorType indVType = getIndexVectorType();
5517	VectorType maskVType = getMaskVectorType();
5518	VectorType valueVType = getVectorType();
5519	MemRefType memType = getMemRefType();
5520
5521	if (valueVType.getElementType() != memType.getElementType())
5522	return emitOpError("base and valueToStore element type should match");
5523	if (llvm::size(getIndices()) != memType.getRank())
5524	return emitOpError("requires ") << memType.getRank() << " indices";
5525	if (valueVType.getShape() != indVType.getShape())
5526	return emitOpError("expected valueToStore dim to match indices dim");
5527	if (valueVType.getShape() != maskVType.getShape())
5528	return emitOpError("expected valueToStore dim to match mask dim");
5529	return success();
5530	}
5531
5532	namespace {
5533	class ScatterFolder final : public OpRewritePattern<ScatterOp> {
5534	public:
5535	using OpRewritePattern::OpRewritePattern;
5536	LogicalResult matchAndRewrite(ScatterOp scatter,
5537	PatternRewriter &rewriter) const override {
5538	switch (getMaskFormat(scatter.getMask())) {
5539	case MaskFormat::AllTrue:
5540	return failure(); // no unmasked equivalent
5541	case MaskFormat::AllFalse:
5542	rewriter.eraseOp(op: scatter);
5543	return success();
5544	case MaskFormat::Unknown:
5545	return failure();
5546	}
5547	llvm_unreachable("Unexpected 1DMaskFormat on ScatterFolder");
5548	}
5549	};
5550
5551	/// Fold scatters with consecutive offsets [0, 1, 2, ...] into contiguous
5552	/// maskedstore. Only 1D fixed vectors are supported for now.
5553	class FoldContiguousScatter final : public OpRewritePattern<ScatterOp> {
5554	public:
5555	using OpRewritePattern::OpRewritePattern;
5556	LogicalResult matchAndRewrite(ScatterOp op,
5557	PatternRewriter &rewriter) const override {
5558	if (failed(isZeroBasedContiguousSeq(op.getIndexVec())))
5559	return failure();
5560
5561	rewriter.replaceOpWithNewOp<MaskedStoreOp>(
5562	op, op.getBase(), op.getIndices(), op.getMask(), op.getValueToStore());
5563	return success();
5564	}
5565	};
5566	} // namespace
5567
5568	void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &results,
5569	MLIRContext *context) {
5570	results.add<ScatterFolder, FoldContiguousScatter>(context);
5571	}
5572
5573	//===----------------------------------------------------------------------===//
5574	// ExpandLoadOp
5575	//===----------------------------------------------------------------------===//
5576
5577	LogicalResult ExpandLoadOp::verify() {
5578	VectorType maskVType = getMaskVectorType();
5579	VectorType passVType = getPassThruVectorType();
5580	VectorType resVType = getVectorType();
5581	MemRefType memType = getMemRefType();
5582
5583	if (resVType.getElementType() != memType.getElementType())
5584	return emitOpError("base and result element type should match");
5585	if (llvm::size(getIndices()) != memType.getRank())
5586	return emitOpError("requires ") << memType.getRank() << " indices";
5587	if (resVType.getDimSize(0) != maskVType.getDimSize(0))
5588	return emitOpError("expected result dim to match mask dim");
5589	if (resVType != passVType)
5590	return emitOpError("expected pass_thru of same type as result type");
5591	return success();
5592	}
5593
5594	namespace {
5595	class ExpandLoadFolder final : public OpRewritePattern<ExpandLoadOp> {
5596	public:
5597	using OpRewritePattern::OpRewritePattern;
5598	LogicalResult matchAndRewrite(ExpandLoadOp expand,
5599	PatternRewriter &rewriter) const override {
5600	switch (getMaskFormat(expand.getMask())) {
5601	case MaskFormat::AllTrue:
5602	rewriter.replaceOpWithNewOp<vector::LoadOp>(
5603	expand, expand.getType(), expand.getBase(), expand.getIndices());
5604	return success();
5605	case MaskFormat::AllFalse:
5606	rewriter.replaceOp(expand, expand.getPassThru());
5607	return success();
5608	case MaskFormat::Unknown:
5609	return failure();
5610	}
5611	llvm_unreachable("Unexpected 1DMaskFormat on ExpandLoadFolder");
5612	}
5613	};
5614	} // namespace
5615
5616	void ExpandLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
5617	MLIRContext *context) {
5618	results.add<ExpandLoadFolder>(context);
5619	}
5620
5621	//===----------------------------------------------------------------------===//
5622	// CompressStoreOp
5623	//===----------------------------------------------------------------------===//
5624
5625	LogicalResult CompressStoreOp::verify() {
5626	VectorType maskVType = getMaskVectorType();
5627	VectorType valueVType = getVectorType();
5628	MemRefType memType = getMemRefType();
5629
5630	if (valueVType.getElementType() != memType.getElementType())
5631	return emitOpError("base and valueToStore element type should match");
5632	if (llvm::size(getIndices()) != memType.getRank())
5633	return emitOpError("requires ") << memType.getRank() << " indices";
5634	if (valueVType.getDimSize(0) != maskVType.getDimSize(0))
5635	return emitOpError("expected valueToStore dim to match mask dim");
5636	return success();
5637	}
5638
5639	namespace {
5640	class CompressStoreFolder final : public OpRewritePattern<CompressStoreOp> {
5641	public:
5642	using OpRewritePattern::OpRewritePattern;
5643	LogicalResult matchAndRewrite(CompressStoreOp compress,
5644	PatternRewriter &rewriter) const override {
5645	switch (getMaskFormat(compress.getMask())) {
5646	case MaskFormat::AllTrue:
5647	rewriter.replaceOpWithNewOp<vector::StoreOp>(
5648	compress, compress.getValueToStore(), compress.getBase(),
5649	compress.getIndices());
5650	return success();
5651	case MaskFormat::AllFalse:
5652	rewriter.eraseOp(op: compress);
5653	return success();
5654	case MaskFormat::Unknown:
5655	return failure();
5656	}
5657	llvm_unreachable("Unexpected 1DMaskFormat on CompressStoreFolder");
5658	}
5659	};
5660	} // namespace
5661
5662	void CompressStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
5663	MLIRContext *context) {
5664	results.add<CompressStoreFolder>(context);
5665	}
5666
5667	//===----------------------------------------------------------------------===//
5668	// ShapeCastOp
5669	//===----------------------------------------------------------------------===//
5670
5671	void ShapeCastOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
5672	SetIntRangeFn setResultRanges) {
5673	setResultRanges(getResult(), argRanges.front());
5674	}
5675
5676	LogicalResult ShapeCastOp::verify() {
5677
5678	VectorType sourceType = getSourceVectorType();
5679	VectorType resultType = getResultVectorType();
5680
5681	// Check that element type is preserved
5682	if (sourceType.getElementType() != resultType.getElementType())
5683	return emitOpError("has different source and result element types");
5684
5685	// Check that number of elements is preserved
5686	int64_t sourceNElms = sourceType.getNumElements();
5687	int64_t resultNElms = resultType.getNumElements();
5688	if (sourceNElms != resultNElms) {
5689	return emitOpError() << "has different number of elements at source ("
5690	<< sourceNElms << ") and result (" << resultNElms
5691	<< ")";
5692	}
5693
5694	// Check that (non-)scalability is preserved
5695	int64_t sourceNScalableDims = sourceType.getNumScalableDims();
5696	int64_t resultNScalableDims = resultType.getNumScalableDims();
5697	if (sourceNScalableDims != resultNScalableDims)
5698	return emitOpError() << "has different number of scalable dims at source ("
5699	<< sourceNScalableDims << ") and result ("
5700	<< resultNScalableDims << ")";
5701
5702	return success();
5703	}
5704
5705	/// Return true if `transpose` does not permute a pair of non-unit dims.
5706	/// By `order preserving` we mean that the flattened versions of the input and
5707	/// output vectors are (numerically) identical. In other words `transpose` is
5708	/// effectively a shape cast.
5709	static bool isOrderPreserving(TransposeOp transpose) {
5710	ArrayRef<int64_t> permutation = transpose.getPermutation();
5711	VectorType sourceType = transpose.getSourceVectorType();
5712	ArrayRef<int64_t> inShape = sourceType.getShape();
5713	ArrayRef<bool> inDimIsScalable = sourceType.getScalableDims();
5714	auto isNonScalableUnitDim = [&](int64_t dim) {
5715	return inShape[dim] == 1 && !inDimIsScalable[dim];
5716	};
5717	int64_t current = 0;
5718	for (auto p : permutation) {
5719	if (!isNonScalableUnitDim(p)) {
5720	if (p < current) {
5721	return false;
5722	}
5723	current = p;
5724	}
5725	}
5726	return true;
5727	}
5728
5729	OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) {
5730
5731	VectorType resultType = getType();
5732
5733	// No-op shape cast.
5734	if (getSource().getType() == resultType)
5735	return getSource();
5736
5737	// shape_cast(shape_cast(x)) -> shape_cast(x)
5738	if (auto precedingShapeCast = getSource().getDefiningOp<ShapeCastOp>()) {
5739	setOperand(precedingShapeCast.getSource());
5740	return getResult();
5741	}
5742
5743	// shape_cast(transpose(x)) -> shape_cast(x)
5744	if (auto transpose = getSource().getDefiningOp<TransposeOp>()) {
5745	// This folder does
5746	// shape_cast(transpose) -> shape_cast
5747	// But another pattern, ConvertIllegalShapeCastOpsToTransposes, does
5748	// shape_cast -> shape_cast(transpose)
5749	// i.e. the complete opposite. When paired, these 2 patterns can cause
5750	// infinite cycles in pattern rewriting.
5751	// ConvertIllegalShapeCastOpsToTransposes only matches on scalable
5752	// vectors, so by disabling this folder for scalable vectors the
5753	// cycle is avoided.
5754	// TODO: Check if ConvertIllegalShapeCastOpsToTransposes is
5755	// still needed. If it's not, then we can fold here.
5756	if (!transpose.getType().isScalable() && isOrderPreserving(transpose)) {
5757	setOperand(transpose.getVector());
5758	return getResult();
5759	}
5760	return {};
5761	}
5762
5763	// Y = shape_cast(broadcast(X))
5764	// -> X, if X and Y have same type
5765	if (auto bcastOp = getSource().getDefiningOp<BroadcastOp>()) {
5766	if (bcastOp.getSourceType() == resultType)
5767	return bcastOp.getSource();
5768	}
5769
5770	// shape_cast(constant) -> constant
5771	if (auto splatAttr =
5772	llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource()))
5773	return splatAttr.reshape(getType());
5774
5775	// shape_cast(poison) -> poison
5776	if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getSource())) {
5777	return ub::PoisonAttr::get(getContext());
5778	}
5779
5780	return {};
5781	}
5782
5783	namespace {
5784
5785	/// Helper function that computes a new vector type based on the input vector
5786	/// type by removing the trailing one dims:
5787	///
5788	/// vector<4x1x1xi1> --> vector<4x1xi1>
5789	///
5790	static VectorType trimTrailingOneDims(VectorType oldType) {
5791	ArrayRef<int64_t> oldShape = oldType.getShape();
5792	ArrayRef<int64_t> newShape = oldShape;
5793
5794	ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
5795	ArrayRef<bool> newScalableDims = oldScalableDims;
5796
5797	while (!newShape.empty() && newShape.back() == 1 && !newScalableDims.back()) {
5798	newShape = newShape.drop_back(N: 1);
5799	newScalableDims = newScalableDims.drop_back(N: 1);
5800	}
5801
5802	// Make sure we have at least 1 dimension.
5803	// TODO: Add support for 0-D vectors.
5804	if (newShape.empty()) {
5805	newShape = oldShape.take_back();
5806	newScalableDims = oldScalableDims.take_back();
5807	}
5808
5809	return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
5810	}
5811
5812	/// Folds qualifying shape_cast(create_mask) into a new create_mask
5813	///
5814	/// Looks at `vector.shape_cast` Ops that simply "drop" the trailing unit
5815	/// dimension. If the input vector comes from `vector.create_mask` for which
5816	/// the corresponding mask input value is 1 (e.g. `%c1` below), then it is safe
5817	/// to fold shape_cast into create_mask.
5818	///
5819	/// BEFORE:
5820	/// %1 = vector.create_mask %c1, %dim, %c1, %c1 : vector<1x[4]x1x1xi1>
5821	/// %2 = vector.shape_cast %1 : vector<1x[4]x1x1xi1> to vector<1x[4]xi1>
5822	/// AFTER:
5823	/// %0 = vector.create_mask %c1, %dim : vector<1x[4]xi1>
5824	class ShapeCastCreateMaskFolderTrailingOneDim final
5825	: public OpRewritePattern<ShapeCastOp> {
5826	public:
5827	using OpRewritePattern::OpRewritePattern;
5828
5829	LogicalResult matchAndRewrite(ShapeCastOp shapeOp,
5830	PatternRewriter &rewriter) const override {
5831	Value shapeOpSrc = shapeOp->getOperand(0);
5832	auto createMaskOp = shapeOpSrc.getDefiningOp<vector::CreateMaskOp>();
5833	auto constantMaskOp = shapeOpSrc.getDefiningOp<vector::ConstantMaskOp>();
5834	if (!createMaskOp && !constantMaskOp)
5835	return failure();
5836
5837	VectorType shapeOpResTy = shapeOp.getResultVectorType();
5838	VectorType shapeOpSrcTy = shapeOp.getSourceVectorType();
5839
5840	VectorType newVecType = trimTrailingOneDims(shapeOpSrcTy);
5841	if (newVecType != shapeOpResTy)
5842	return failure();
5843
5844	auto numDimsToDrop =
5845	shapeOpSrcTy.getShape().size() - shapeOpResTy.getShape().size();
5846
5847	// No unit dims to drop
5848	if (!numDimsToDrop)
5849	return failure();
5850
5851	if (createMaskOp) {
5852	auto maskOperands = createMaskOp.getOperands();
5853	auto numMaskOperands = maskOperands.size();
5854
5855	// Check every mask dim size to see whether it can be dropped
5856	for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
5857	--i) {
5858	auto constant = maskOperands[i].getDefiningOp<arith::ConstantIndexOp>();
5859	if (!constant || (constant.value() != 1))
5860	return failure();
5861	}
5862	SmallVector<Value> newMaskOperands =
5863	maskOperands.drop_back(numDimsToDrop);
5864
5865	rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(shapeOp, shapeOpResTy,
5866	newMaskOperands);
5867	return success();
5868	}
5869
5870	if (constantMaskOp) {
5871	auto maskDimSizes = constantMaskOp.getMaskDimSizes();
5872	auto numMaskOperands = maskDimSizes.size();
5873
5874	// Check every mask dim size to see whether it can be dropped
5875	for (size_t i = numMaskOperands - 1; i >= numMaskOperands - numDimsToDrop;
5876	--i) {
5877	if (maskDimSizes[i] != 1)
5878	return failure();
5879	}
5880
5881	auto newMaskOperands = maskDimSizes.drop_back(numDimsToDrop);
5882	rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(shapeOp, shapeOpResTy,
5883	newMaskOperands);
5884	return success();
5885	}
5886
5887	return failure();
5888	}
5889	};
5890
5891	/// Pattern to rewrite Y = ShapeCast(Broadcast(X)) as either
5892	/// i) Y = ShapeCast(X), or
5893	/// ii) Y = Broadcast(X)
5894	/// If both (i) and (ii) are possible, (i) is chosen.
5895	class ShapeCastBroadcastFolder final : public OpRewritePattern<ShapeCastOp> {
5896	public:
5897	using OpRewritePattern::OpRewritePattern;
5898
5899	LogicalResult matchAndRewrite(ShapeCastOp shapeCastOp,
5900	PatternRewriter &rewriter) const override {
5901	auto broadcastOp =
5902	shapeCastOp.getSource().getDefiningOp<vector::BroadcastOp>();
5903	if (!broadcastOp)
5904	return failure();
5905
5906	auto srcVectorType = dyn_cast<VectorType>(broadcastOp.getSourceType());
5907	bool srcIsScalar = !srcVectorType;
5908
5909	// Replace Y = ShapeCast(Broadcast(X)) with Y = ShapeCast(X).
5910	// Example:
5911	// %0 = vector.broadcast %in : vector<3x4xf32> to vector<1x3x4xf32>
5912	// %1 = vector.shape_cast %0 : vector<1x3x4xf32> to vector<12xf32>
5913	// to
5914	// %1 = vector.shape_cast %in : vector<3x4xf32> to vector<12xf32>
5915	if (srcVectorType) {
5916	if (srcVectorType.getNumElements() ==
5917	shapeCastOp.getResultVectorType().getNumElements()) {
5918	rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
5919	shapeCastOp, shapeCastOp.getResultVectorType(),
5920	broadcastOp.getSource());
5921	return success();
5922	}
5923	}
5924
5925	// Replace Y = ShapeCast(Broadcast(X)) with Y = Broadcast(X)
5926	// Example
5927	// %0 = vector.broadcast %in : vector<3xf32> to vector<2x4x3xf32>
5928	// %1 = vector.shape_cast %0 : vector<2x4x3xf32> to vector<8x3xf32>
5929	// to
5930	// %1 = vector.broadcast %in : vector<3xf32> to vector<8x3xf32>
5931	VectorType dstVectorType = shapeCastOp.getResultVectorType();
5932	if (srcIsScalar || isBroadcastableTo(srcVectorType, dstVectorType) ==
5933	BroadcastableToResult::Success) {
5934	rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
5935	shapeCastOp, dstVectorType, broadcastOp.getSource());
5936	return success();
5937	}
5938	return failure();
5939	}
5940	};
5941
5942	} // namespace
5943
5944	void ShapeCastOp::getCanonicalizationPatterns(RewritePatternSet &results,
5945	MLIRContext *context) {
5946	results
5947	.add<ShapeCastCreateMaskFolderTrailingOneDim, ShapeCastBroadcastFolder>(
5948	context);
5949	}
5950
5951	//===----------------------------------------------------------------------===//
5952	// VectorBitCastOp
5953	//===----------------------------------------------------------------------===//
5954
5955	LogicalResult BitCastOp::verify() {
5956	auto sourceVectorType = getSourceVectorType();
5957	auto resultVectorType = getResultVectorType();
5958
5959	for (int64_t i = 0, e = sourceVectorType.getRank() - 1; i < e; i++) {
5960	if (sourceVectorType.getDimSize(i) != resultVectorType.getDimSize(i))
5961	return emitOpError("dimension size mismatch at: ") << i;
5962	}
5963
5964	DataLayout dataLayout = DataLayout::closest(*this);
5965	auto sourceElementBits =
5966	dataLayout.getTypeSizeInBits(sourceVectorType.getElementType());
5967	auto resultElementBits =
5968	dataLayout.getTypeSizeInBits(resultVectorType.getElementType());
5969
5970	if (sourceVectorType.getRank() == 0) {
5971	if (sourceElementBits != resultElementBits)
5972	return emitOpError("source/result bitwidth of the 0-D vector element "
5973	"types must be equal");
5974	} else if (sourceElementBits * sourceVectorType.getShape().back() !=
5975	resultElementBits * resultVectorType.getShape().back()) {
5976	return emitOpError(
5977	"source/result bitwidth of the minor 1-D vectors must be equal");
5978	}
5979
5980	return success();
5981	}
5982
5983	OpFoldResult BitCastOp::fold(FoldAdaptor adaptor) {
5984	// Nop cast.
5985	if (getSource().getType() == getResult().getType())
5986	return getSource();
5987
5988	// Canceling bitcasts.
5989	if (auto otherOp = getSource().getDefiningOp<BitCastOp>()) {
5990	if (getResult().getType() == otherOp.getSource().getType())
5991	return otherOp.getSource();
5992
5993	setOperand(otherOp.getSource());
5994	return getResult();
5995	}
5996
5997	Attribute sourceConstant = adaptor.getSource();
5998	if (!sourceConstant)
5999	return {};
6000
6001	Type srcElemType = getSourceVectorType().getElementType();
6002	Type dstElemType = getResultVectorType().getElementType();
6003
6004	if (auto floatPack = llvm::dyn_cast<DenseFPElementsAttr>(sourceConstant)) {
6005	if (floatPack.isSplat()) {
6006	auto splat = floatPack.getSplatValue<FloatAttr>();
6007
6008	// Casting fp16 into fp32.
6009	if (srcElemType.isF16() && dstElemType.isF32()) {
6010	uint32_t bits = static_cast<uint32_t>(
6011	splat.getValue().bitcastToAPInt().getZExtValue());
6012	// Duplicate the 16-bit pattern.
6013	bits = (bits << 16) | (bits & 0xffff);
6014	APInt intBits(32, bits);
6015	APFloat floatBits(llvm::APFloat::IEEEsingle(), intBits);
6016	return DenseElementsAttr::get(getResultVectorType(), floatBits);
6017	}
6018	}
6019	}
6020
6021	if (auto intPack = llvm::dyn_cast<DenseIntElementsAttr>(sourceConstant)) {
6022	if (intPack.isSplat()) {
6023	auto splat = intPack.getSplatValue<IntegerAttr>();
6024
6025	if (llvm::isa<IntegerType>(dstElemType)) {
6026	uint64_t srcBitWidth = srcElemType.getIntOrFloatBitWidth();
6027	uint64_t dstBitWidth = dstElemType.getIntOrFloatBitWidth();
6028
6029	// Casting to a larger integer bit width.
6030	if (dstBitWidth > srcBitWidth && dstBitWidth % srcBitWidth == 0) {
6031	APInt intBits = splat.getValue().zext(dstBitWidth);
6032
6033	// Duplicate the lower width element.
6034	for (uint64_t i = 0; i < dstBitWidth / srcBitWidth - 1; i++)
6035	intBits = (intBits << srcBitWidth) | intBits;
6036	return DenseElementsAttr::get(getResultVectorType(), intBits);
6037	}
6038	}
6039	}
6040	}
6041
6042	return {};
6043	}
6044
6045	//===----------------------------------------------------------------------===//
6046	// TypeCastOp
6047	//===----------------------------------------------------------------------===//
6048
6049	static SmallVector<int64_t, 8> extractShape(MemRefType memRefType) {
6050	auto vectorType = llvm::dyn_cast<VectorType>(memRefType.getElementType());
6051	SmallVector<int64_t, 8> res(memRefType.getShape());
6052	if (vectorType)
6053	res.append(vectorType.getShape().begin(), vectorType.getShape().end());
6054	return res;
6055	}
6056
6057	/// Build the canonical memRefType with a single vector.
6058	/// E.g. memref<4 x 5 x vector<6 x f32>> -> memref<vector<4 x 5 x 6 x f32>>.
6059	void TypeCastOp::build(OpBuilder &builder, OperationState &result,
6060	Value source) {
6061	result.addOperands(source);
6062	MemRefType memRefType = llvm::cast<MemRefType>(source.getType());
6063	VectorType vectorType =
6064	VectorType::get(extractShape(memRefType),
6065	getElementTypeOrSelf(getElementTypeOrSelf(memRefType)));
6066	result.addTypes(MemRefType::get({}, vectorType, MemRefLayoutAttrInterface(),
6067	memRefType.getMemorySpace()));
6068	}
6069
6070	LogicalResult TypeCastOp::verify() {
6071	MemRefType canonicalType = getMemRefType().canonicalizeStridedLayout();
6072	if (!canonicalType.getLayout().isIdentity())
6073	return emitOpError("expects operand to be a memref with identity layout");
6074	if (!getResultMemRefType().getLayout().isIdentity())
6075	return emitOpError("expects result to be a memref with identity layout");
6076	if (getResultMemRefType().getMemorySpace() !=
6077	getMemRefType().getMemorySpace())
6078	return emitOpError("expects result in same memory space");
6079
6080	auto sourceType = getMemRefType();
6081	auto resultType = getResultMemRefType();
6082	if (getElementTypeOrSelf(getElementTypeOrSelf(sourceType)) !=
6083	getElementTypeOrSelf(getElementTypeOrSelf(resultType)))
6084	return emitOpError(
6085	"expects result and operand with same underlying scalar type: ")
6086	<< resultType;
6087	if (extractShape(sourceType) != extractShape(resultType))
6088	return emitOpError(
6089	"expects concatenated result and operand shapes to be equal: ")
6090	<< resultType;
6091	return success();
6092	}
6093
6094	//===----------------------------------------------------------------------===//
6095	// TransposeOp
6096	//===----------------------------------------------------------------------===//
6097
6098	void vector::TransposeOp::build(OpBuilder &builder, OperationState &result,
6099	Value vector, ArrayRef<int64_t> permutation) {
6100	VectorType vt = llvm::cast<VectorType>(vector.getType());
6101	SmallVector<int64_t, 4> transposedShape(vt.getRank());
6102	SmallVector<bool, 4> transposedScalableDims(vt.getRank());
6103	for (unsigned i = 0; i < permutation.size(); ++i) {
6104	transposedShape[i] = vt.getShape()[permutation[i]];
6105	transposedScalableDims[i] = vt.getScalableDims()[permutation[i]];
6106	}
6107
6108	result.addOperands(vector);
6109	result.addTypes(VectorType::get(transposedShape, vt.getElementType(),
6110	transposedScalableDims));
6111	result.addAttribute(TransposeOp::getPermutationAttrName(result.name),
6112	builder.getDenseI64ArrayAttr(permutation));
6113	}
6114
6115	OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
6116	// Eliminate splat constant transpose ops.
6117	if (auto splat =
6118	llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getVector()))
6119	return splat.reshape(getResultVectorType());
6120
6121	// Eliminate poison transpose ops.
6122	if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getVector()))
6123	return ub::PoisonAttr::get(getContext());
6124
6125	// Eliminate identity transposes, and more generally any transposes that
6126	// preserves the shape without permuting elements.
6127	//
6128	// Examples of what to fold:
6129	// %0 = vector.transpose %arg, [0, 1] : vector<1x1xi8> to vector<1x1xi8>
6130	// %0 = vector.transpose %arg, [0, 1] : vector<2x2xi8> to vector<2x2xi8>
6131	// %0 = vector.transpose %arg, [1, 0] : vector<1x1xi8> to vector<1x1xi8>
6132	//
6133	// Example of what NOT to fold:
6134	// %0 = vector.transpose %arg, [1, 0] : vector<2x2xi8> to vector<2x2xi8>
6135	//
6136	if (getSourceVectorType() == getResultVectorType() &&
6137	isOrderPreserving(*this))
6138	return getVector();
6139
6140	return {};
6141	}
6142
6143	LogicalResult vector::TransposeOp::verify() {
6144	VectorType vectorType = getSourceVectorType();
6145	VectorType resultType = getResultVectorType();
6146	int64_t rank = resultType.getRank();
6147	if (vectorType.getRank() != rank)
6148	return emitOpError("vector result rank mismatch: ") << rank;
6149	// Verify transposition array.
6150	ArrayRef<int64_t> perm = getPermutation();
6151	int64_t size = perm.size();
6152	if (rank != size)
6153	return emitOpError("transposition length mismatch: ") << size;
6154	SmallVector<bool, 8> seen(rank, false);
6155	for (const auto &ta : llvm::enumerate(perm)) {
6156	if (ta.value() < 0 || ta.value() >= rank)
6157	return emitOpError("transposition index out of range: ") << ta.value();
6158	if (seen[ta.value()])
6159	return emitOpError("duplicate position index: ") << ta.value();
6160	seen[ta.value()] = true;
6161	if (resultType.getDimSize(ta.index()) != vectorType.getDimSize(ta.value()))
6162	return emitOpError("dimension size mismatch at: ") << ta.value();
6163	}
6164	return success();
6165	}
6166
6167	std::optional<SmallVector<int64_t, 4>> TransposeOp::getShapeForUnroll() {
6168	return llvm::to_vector<4>(getResultVectorType().getShape());
6169	}
6170
6171	namespace {
6172
6173	// Rewrites two back-to-back TransposeOp operations into a single TransposeOp.
6174	class TransposeFolder final : public OpRewritePattern<vector::TransposeOp> {
6175	public:
6176	using OpRewritePattern::OpRewritePattern;
6177
6178	LogicalResult matchAndRewrite(vector::TransposeOp transposeOp,
6179	PatternRewriter &rewriter) const override {
6180	// Composes two permutations: result[i] = permutation1[permutation2[i]].
6181	auto composePermutations = [](ArrayRef<int64_t> permutation1,
6182	ArrayRef<int64_t> permutation2) {
6183	SmallVector<int64_t, 4> result;
6184	for (auto index : permutation2)
6185	result.push_back(Elt: permutation1[index]);
6186	return result;
6187	};
6188
6189	// Return if the input of 'transposeOp' is not defined by another transpose.
6190	vector::TransposeOp parentTransposeOp =
6191	transposeOp.getVector().getDefiningOp<vector::TransposeOp>();
6192	if (!parentTransposeOp)
6193	return failure();
6194
6195	SmallVector<int64_t, 4> permutation = composePermutations(
6196	parentTransposeOp.getPermutation(), transposeOp.getPermutation());
6197	// Replace 'transposeOp' with a new transpose operation.
6198	rewriter.replaceOpWithNewOp<vector::TransposeOp>(
6199	transposeOp, transposeOp.getResult().getType(),
6200	parentTransposeOp.getVector(), permutation);
6201	return success();
6202	}
6203	};
6204
6205	// Folds transpose(splat x : src_type) : res_type into splat x : res_type.
6206	class FoldTransposeSplat final : public OpRewritePattern<TransposeOp> {
6207	public:
6208	using OpRewritePattern::OpRewritePattern;
6209
6210	LogicalResult matchAndRewrite(TransposeOp transposeOp,
6211	PatternRewriter &rewriter) const override {
6212	auto splatOp = transposeOp.getVector().getDefiningOp<vector::SplatOp>();
6213	if (!splatOp)
6214	return failure();
6215
6216	rewriter.replaceOpWithNewOp<vector::SplatOp>(
6217	transposeOp, transposeOp.getResultVectorType(), splatOp.getInput());
6218	return success();
6219	}
6220	};
6221
6222	/// Folds transpose(create_mask) into a new transposed create_mask.
6223	class FoldTransposeCreateMask final : public OpRewritePattern<TransposeOp> {
6224	public:
6225	using OpRewritePattern::OpRewritePattern;
6226
6227	LogicalResult matchAndRewrite(TransposeOp transpOp,
6228	PatternRewriter &rewriter) const override {
6229	Value transposeSrc = transpOp.getVector();
6230	auto createMaskOp = transposeSrc.getDefiningOp<vector::CreateMaskOp>();
6231	auto constantMaskOp = transposeSrc.getDefiningOp<vector::ConstantMaskOp>();
6232	if (!createMaskOp && !constantMaskOp)
6233	return failure();
6234
6235	// Get the transpose permutation and apply it to the vector.create_mask or
6236	// vector.constant_mask operands.
6237	ArrayRef<int64_t> permutation = transpOp.getPermutation();
6238
6239	if (createMaskOp) {
6240	auto maskOperands = createMaskOp.getOperands();
6241	SmallVector<Value> newOperands(maskOperands.begin(), maskOperands.end());
6242	applyPermutationToVector(inVec&: newOperands, permutation);
6243
6244	rewriter.replaceOpWithNewOp<vector::CreateMaskOp>(
6245	transpOp, transpOp.getResultVectorType(), newOperands);
6246	return success();
6247	}
6248
6249	// ConstantMaskOp case.
6250	auto maskDimSizes = constantMaskOp.getMaskDimSizes();
6251	auto newMaskDimSizes = applyPermutation(maskDimSizes, permutation);
6252
6253	rewriter.replaceOpWithNewOp<vector::ConstantMaskOp>(
6254	transpOp, transpOp.getResultVectorType(), newMaskDimSizes);
6255	return success();
6256	}
6257	};
6258
6259	/// Folds transpose(shape_cast) into a new shape_cast.
6260	class FoldTransposeShapeCast final : public OpRewritePattern<TransposeOp> {
6261	public:
6262	using OpRewritePattern::OpRewritePattern;
6263
6264	LogicalResult matchAndRewrite(TransposeOp transposeOp,
6265	PatternRewriter &rewriter) const override {
6266	auto shapeCastOp =
6267	transposeOp.getVector().getDefiningOp<vector::ShapeCastOp>();
6268	if (!shapeCastOp)
6269	return failure();
6270	if (!isOrderPreserving(transposeOp))
6271	return failure();
6272
6273	VectorType resultType = transposeOp.getType();
6274
6275	// We don't need to check isValidShapeCast at this point, because it is
6276	// guaranteed that merging the transpose into the the shape_cast is a valid
6277	// shape_cast, because the transpose just inserts/removes ones.
6278
6279	rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(transposeOp, resultType,
6280	shapeCastOp.getSource());
6281	return success();
6282	}
6283	};
6284
6285	/// Folds transpose(broadcast(x)) to broadcast(x) if the transpose is
6286	/// 'order preserving', where 'order preserving' means the flattened
6287	/// inputs and outputs of the transpose have identical (numerical) values.
6288	///
6289	/// Example:
6290	/// ```
6291	/// %0 = vector.broadcast %input : vector<1x1xi32> to vector<1x8xi32>
6292	/// %1 = vector.transpose %0, [1, 0] : vector<1x8xi32>
6293	/// to vector<8x1xi32>
6294	/// ```
6295	/// can be rewritten as the equivalent
6296	/// ```
6297	/// %0 = vector.broadcast %input : vector<1x1xi32> to vector<8x1xi32>.
6298	/// ```
6299	/// The algorithm works by partitioning dimensions into groups that can be
6300	/// locally permuted while preserving order, and checks that the transpose
6301	/// only permutes within these groups.
6302	///
6303	/// Groups are either contiguous sequences of 1s, or non-1s (1-element groups).
6304	/// Consider broadcasting 4x1x1x7 to 2x3x4x5x6x7. This is equivalent to
6305	/// broadcasting from 1x1x4x1x1x7.
6306	/// ^^^ ^ ^^^ ^
6307	/// groups: 0 1 2 3
6308	/// Order preserving permutations for this example are ones that only permute
6309	/// within the groups [0,1] and [3,4], like (1 0 2 4 3 5 6).
6310	class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
6311	public:
6312	using OpRewritePattern::OpRewritePattern;
6313	FoldTransposeBroadcast(MLIRContext *context, PatternBenefit benefit = 1)
6314	: OpRewritePattern<vector::TransposeOp>(context, benefit) {}
6315
6316	LogicalResult matchAndRewrite(vector::TransposeOp transpose,
6317	PatternRewriter &rewriter) const override {
6318
6319	vector::BroadcastOp broadcast =
6320	transpose.getVector().getDefiningOp<vector::BroadcastOp>();
6321	if (!broadcast) {
6322	return rewriter.notifyMatchFailure(transpose,
6323	"not preceded by a broadcast");
6324	}
6325
6326	auto inputType = dyn_cast<VectorType>(broadcast.getSourceType());
6327	VectorType outputType = transpose.getResultVectorType();
6328
6329	// transpose(broadcast(scalar)) -> broadcast(scalar) is always valid
6330	bool inputIsScalar = !inputType;
6331	if (inputIsScalar) {
6332	rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
6333	broadcast.getSource());
6334	return success();
6335	}
6336
6337	ArrayRef<int64_t> permutation = transpose.getPermutation();
6338	ArrayRef<int64_t> inputShape = inputType.getShape();
6339	int64_t inputRank = inputType.getRank();
6340	int64_t outputRank = transpose.getType().getRank();
6341	int64_t deltaRank = outputRank - inputRank;
6342
6343	int low = 0;
6344	for (int inputIndex = 0; inputIndex < inputRank; ++inputIndex) {
6345	bool notOne = inputShape[inputIndex] != 1;
6346	bool prevNotOne = (inputIndex != 0 && inputShape[inputIndex - 1] != 1);
6347	bool groupEndFound = notOne || prevNotOne;
6348	if (groupEndFound) {
6349	int high = inputIndex + deltaRank;
6350	// Return failure if not all permutation destinations for indices in
6351	// [low, high) are in [low, high), i.e. the permutation is not local to
6352	// the group.
6353	for (int i = low; i < high; ++i) {
6354	if (permutation[i] < low || permutation[i] >= high) {
6355	return rewriter.notifyMatchFailure(
6356	transpose, "permutation not local to group");
6357	}
6358	}
6359	low = high;
6360	}
6361	}
6362
6363	// We don't need to check the final group [low, outputRank) because if it is
6364	// not locally bound, there must be a preceding group that already failed
6365	// the check (impossible to have just 1 non-locally bound group).
6366
6367	// The preceding logic also ensures that at this point, the output of the
6368	// transpose is definitely broadcastable from the input shape, assert so:
6369	assert(vector::isBroadcastableTo(inputType, outputType) ==
6370	vector::BroadcastableToResult::Success &&
6371	"not broadcastable directly to transpose output");
6372
6373	rewriter.replaceOpWithNewOp<vector::BroadcastOp>(transpose, outputType,
6374	broadcast.getSource());
6375
6376	return success();
6377	}
6378	};
6379
6380	} // namespace
6381
6382	void vector::TransposeOp::getCanonicalizationPatterns(
6383	RewritePatternSet &results, MLIRContext *context) {
6384	results.add<FoldTransposeCreateMask, FoldTransposeShapeCast, TransposeFolder,
6385	FoldTransposeSplat, FoldTransposeBroadcast>(context);
6386	}
6387
6388	//===----------------------------------------------------------------------===//
6389	// ConstantMaskOp
6390	//===----------------------------------------------------------------------===//
6391
6392	void ConstantMaskOp::build(OpBuilder &builder, OperationState &result,
6393	VectorType type, ConstantMaskKind kind) {
6394	assert(kind == ConstantMaskKind::AllTrue ||
6395	kind == ConstantMaskKind::AllFalse);
6396	build(builder, result, type,
6397	kind == ConstantMaskKind::AllTrue
6398	? type.getShape()
6399	: SmallVector<int64_t>(type.getRank(), 0));
6400	}
6401
6402	LogicalResult ConstantMaskOp::verify() {
6403	auto resultType = llvm::cast<VectorType>(getResult().getType());
6404	// Check the corner case of 0-D vectors first.
6405	if (resultType.getRank() == 0) {
6406	if (getMaskDimSizes().size() != 1)
6407	return emitError("array attr must have length 1 for 0-D vectors");
6408	auto dim = getMaskDimSizes()[0];
6409	if (dim != 0 && dim != 1)
6410	return emitError("mask dim size must be either 0 or 1 for 0-D vectors");
6411	return success();
6412	}
6413
6414	// Verify that array attr size matches the rank of the vector result.
6415	if (static_cast<int64_t>(getMaskDimSizes().size()) != resultType.getRank())
6416	return emitOpError(
6417	"must specify array attr of size equal vector result rank");
6418	// Verify that each array attr element is in bounds of corresponding vector
6419	// result dimension size.
6420	auto resultShape = resultType.getShape();
6421	auto resultScalableDims = resultType.getScalableDims();
6422	ArrayRef<int64_t> maskDimSizes = getMaskDimSizes();
6423	for (const auto [index, maskDimSize] : llvm::enumerate(maskDimSizes)) {
6424	if (maskDimSize < 0 || maskDimSize > resultShape[index])
6425	return emitOpError(
6426	"array attr of size out of bounds of vector result dimension size");
6427	if (resultScalableDims[index] && maskDimSize != 0 &&
6428	maskDimSize != resultShape[index])
6429	return emitOpError(
6430	"only supports 'none set' or 'all set' scalable dimensions");
6431	}
6432	// Verify that if one mask dim size is zero, they all should be zero (because
6433	// the mask region is a conjunction of each mask dimension interval).
6434	bool anyZeros = llvm::is_contained(maskDimSizes, 0);
6435	bool allZeros = llvm::all_of(maskDimSizes, [](int64_t s) { return s == 0; });
6436	if (anyZeros && !allZeros)
6437	return emitOpError("expected all mask dim sizes to be zeros, "
6438	"as a result of conjunction with zero mask dim");
6439	return success();
6440	}
6441
6442	bool ConstantMaskOp::isAllOnesMask() {
6443	auto resultType = getVectorType();
6444	// Check the corner case of 0-D vectors first.
6445	if (resultType.getRank() == 0) {
6446	assert(getMaskDimSizes().size() == 1 && "invalid sizes for zero rank mask");
6447	return getMaskDimSizes()[0] == 1;
6448	}
6449	for (const auto [resultSize, maskDimSize] :
6450	llvm::zip_equal(resultType.getShape(), getMaskDimSizes())) {
6451	if (maskDimSize < resultSize)
6452	return false;
6453	}
6454	return true;
6455	}
6456
6457	//===----------------------------------------------------------------------===//
6458	// CreateMaskOp
6459	//===----------------------------------------------------------------------===//
6460
6461	void CreateMaskOp::build(OpBuilder &builder, OperationState &result,
6462	VectorType type,
6463	ArrayRef<OpFoldResult> mixedOperands) {
6464	SmallVector<Value> operands =
6465	getValueOrCreateConstantIndexOp(builder, result.location, mixedOperands);
6466	build(builder, result, type, operands);
6467	}
6468
6469	LogicalResult CreateMaskOp::verify() {
6470	auto vectorType = llvm::cast<VectorType>(getResult().getType());
6471	// Verify that an operand was specified for each result vector each dimension.
6472	if (vectorType.getRank() == 0) {
6473	if (getNumOperands() != 1)
6474	return emitOpError(
6475	"must specify exactly one operand for 0-D create_mask");
6476	} else if (getNumOperands() !=
6477	llvm::cast<VectorType>(getResult().getType()).getRank()) {
6478	return emitOpError(
6479	"must specify an operand for each result vector dimension");
6480	}
6481	return success();
6482	}
6483
6484	namespace {
6485
6486	/// Pattern to rewrite a CreateMaskOp with a ConstantMaskOp.
6487	///
6488	/// Ex 1:
6489	/// %c2 = arith.constant 2 : index
6490	/// %c3 = arith.constant 3 : index
6491	/// %0 = vector.create_mask %c3, %c2 : vector<4x3xi1>
6492	/// Becomes:
6493	/// vector.constant_mask [3, 2] : vector<4x3xi1>
6494	///
6495	/// Ex 2:
6496	/// %c_neg_1 = arith.constant -1 : index
6497	/// %0 = vector.create_mask %c_neg_1 : vector<[8]xi1>
6498	/// becomes:
6499	/// vector.constant_mask [0] : vector<[8]xi1>
6500	///
6501	/// Ex 3:
6502	/// %c8 = arith.constant 8 : index
6503	/// %c16 = arith.constant 16 : index
6504	/// %0 = vector.vscale
6505	/// %1 = arith.muli %0, %c16 : index
6506	/// %10 = vector.create_mask %c8, %1 : vector<8x[16]xi1>
6507	/// becomes:
6508	/// %0 = vector.constant_mask [8, 16] : vector<8x[16]xi1>
6509	class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
6510	public:
6511	using OpRewritePattern::OpRewritePattern;
6512
6513	LogicalResult matchAndRewrite(CreateMaskOp createMaskOp,
6514	PatternRewriter &rewriter) const override {
6515	VectorType maskType = createMaskOp.getVectorType();
6516	ArrayRef<int64_t> maskTypeDimSizes = maskType.getShape();
6517	ArrayRef<bool> maskTypeDimScalableFlags = maskType.getScalableDims();
6518
6519	// Special case: Rank zero shape.
6520	constexpr std::array<int64_t, 1> rankZeroShape{1};
6521	constexpr std::array<bool, 1> rankZeroScalableDims{false};
6522	if (maskType.getRank() == 0) {
6523	maskTypeDimSizes = rankZeroShape;
6524	maskTypeDimScalableFlags = rankZeroScalableDims;
6525	}
6526
6527	// Determine if this CreateMaskOp can be folded to a ConstantMaskOp and
6528	// collect the `constantDims` (for the ConstantMaskOp).
6529	SmallVector<int64_t, 4> constantDims;
6530	for (auto [i, dimSize] : llvm::enumerate(createMaskOp.getOperands())) {
6531	if (auto intSize = getConstantIntValue(dimSize)) {
6532	// Constant value.
6533	// If the mask dim is non-scalable this can be any value.
6534	// If the mask dim is scalable only zero (all-false) is supported.
6535	if (maskTypeDimScalableFlags[i] && intSize >= 0)
6536	return failure();
6537	constantDims.push_back(*intSize);
6538	} else if (auto vscaleMultiplier = getConstantVscaleMultiplier(dimSize)) {
6539	// Constant vscale multiple (e.g. 4 x vscale).
6540	// Must be all-true to fold to a ConstantMask.
6541	if (vscaleMultiplier < maskTypeDimSizes[i])
6542	return failure();
6543	constantDims.push_back(*vscaleMultiplier);
6544	} else {
6545	return failure();
6546	}
6547	}
6548
6549	// Clamp values to constant_mask bounds.
6550	for (auto [value, maskDimSize] : llvm::zip(constantDims, maskTypeDimSizes))
6551	value = std::clamp<int64_t>(value, 0, maskDimSize);
6552
6553	// If one of dim sizes is zero, set all dims to zero.
6554	if (llvm::is_contained(Range&: constantDims, Element: 0))
6555	constantDims.assign(NumElts: constantDims.size(), Elt: 0);
6556
6557	// Replace 'createMaskOp' with ConstantMaskOp.
6558	rewriter.replaceOpWithNewOp<ConstantMaskOp>(createMaskOp, maskType,
6559	constantDims);
6560	return success();
6561	}
6562	};
6563
6564	} // namespace
6565
6566	void CreateMaskOp::getCanonicalizationPatterns(RewritePatternSet &results,
6567	MLIRContext *context) {
6568	results.add<CreateMaskFolder>(context);
6569	}
6570
6571	//===----------------------------------------------------------------------===//
6572	// MaskOp
6573	//===----------------------------------------------------------------------===//
6574
6575	void MaskOp::build(
6576	OpBuilder &builder, OperationState &result, Value mask,
6577	Operation *maskableOp,
6578	function_ref<void(OpBuilder &, Operation *)> maskRegionBuilder) {
6579	assert(maskRegionBuilder &&
6580	"builder callback for 'maskRegion' must be present");
6581
6582	result.addOperands(mask);
6583	OpBuilder::InsertionGuard guard(builder);
6584	Region *maskRegion = result.addRegion();
6585	builder.createBlock(maskRegion);
6586	maskRegionBuilder(builder, maskableOp);
6587	}
6588
6589	void MaskOp::build(
6590	OpBuilder &builder, OperationState &result, TypeRange resultTypes,
6591	Value mask, Operation *maskableOp,
6592	function_ref<void(OpBuilder &, Operation *)> maskRegionBuilder) {
6593	build(builder, result, resultTypes, mask, /*passthru=*/Value(), maskableOp,
6594	maskRegionBuilder);
6595	}
6596
6597	void MaskOp::build(
6598	OpBuilder &builder, OperationState &result, TypeRange resultTypes,
6599	Value mask, Value passthru, Operation *maskableOp,
6600	function_ref<void(OpBuilder &, Operation *)> maskRegionBuilder) {
6601	build(builder, result, mask, maskableOp, maskRegionBuilder);
6602	if (passthru)
6603	result.addOperands(passthru);
6604	result.addTypes(resultTypes);
6605	}
6606
6607	ParseResult MaskOp::parse(OpAsmParser &parser, OperationState &result) {
6608	// Create the op region.
6609	result.regions.reserve(1);
6610	Region &maskRegion = *result.addRegion();
6611
6612	auto &builder = parser.getBuilder();
6613
6614	// Parse all the operands.
6615	OpAsmParser::UnresolvedOperand mask;
6616	if (parser.parseOperand(mask))
6617	return failure();
6618
6619	// Optional passthru operand.
6620	OpAsmParser::UnresolvedOperand passthru;
6621	ParseResult parsePassthru = parser.parseOptionalComma();
6622	if (parsePassthru.succeeded() && parser.parseOperand(passthru))
6623	return failure();
6624
6625	// Parse op region.
6626	if (parser.parseRegion(maskRegion, /*arguments=*/{}, /*argTypes=*/{}))
6627	return failure();
6628
6629	MaskOp::ensureTerminator(maskRegion, builder, result.location);
6630
6631	// Parse the optional attribute list.
6632	if (parser.parseOptionalAttrDict(result.attributes))
6633	return failure();
6634
6635	// Parse all the types.
6636	Type maskType;
6637	if (parser.parseColonType(maskType))
6638	return failure();
6639
6640	SmallVector<Type> resultTypes;
6641	if (parser.parseOptionalArrowTypeList(resultTypes))
6642	return failure();
6643	result.types.append(resultTypes);
6644
6645	// Resolve operands.
6646	if (parser.resolveOperand(mask, maskType, result.operands))
6647	return failure();
6648
6649	if (parsePassthru.succeeded()) {
6650	if (resultTypes.empty())
6651	return parser.emitError(
6652	parser.getNameLoc(),
6653	"expects a result if passthru operand is provided");
6654
6655	if (parser.resolveOperand(passthru, resultTypes[0], result.operands))
6656	return failure();
6657	}
6658
6659	return success();
6660	}
6661
6662	void mlir::vector::MaskOp::print(OpAsmPrinter &p) {
6663	p << " " << getMask();
6664	if (getPassthru())
6665	p << ", " << getPassthru();
6666
6667	// Print single masked operation and skip terminator.
6668	p << " { ";
6669	Block *singleBlock = &getMaskRegion().getBlocks().front();
6670	if (singleBlock && !singleBlock->getOperations().empty())
6671	p.printCustomOrGenericOp(&singleBlock->front());
6672	p << " }";
6673
6674	p.printOptionalAttrDict(getOperation()->getAttrs());
6675
6676	p << " : " << getMask().getType();
6677	if (getNumResults() > 0)
6678	p << " -> " << getResultTypes();
6679	}
6680
6681	void MaskOp::ensureTerminator(Region &region, Builder &builder, Location loc) {
6682	// 1. For an empty `vector.mask`, create a default terminator.
6683	if (region.empty() || region.front().empty()) {
6684	OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
6685	MaskOp>::ensureTerminator(region, builder, loc);
6686	return;
6687	}
6688
6689	// 2. For a non-empty `vector.mask` with an explicit terminator, do nothing.
6690	Block &block = region.front();
6691	if (isa<vector::YieldOp>(block.back()))
6692	return;
6693
6694	// 3. For a non-empty `vector.mask` without an explicit terminator:
6695
6696	// Create default terminator if the number of masked operations is not
6697	// one. This case will trigger a verification failure.
6698	if (block.getOperations().size() != 1) {
6699	OpTrait::SingleBlockImplicitTerminator<vector::YieldOp>::Impl<
6700	MaskOp>::ensureTerminator(region, builder, loc);
6701	return;
6702	}
6703
6704	// Create a terminator that yields the results from the masked operation.
6705	OpBuilder opBuilder(builder.getContext());
6706	Operation *maskedOp = &block.front();
6707	opBuilder.setInsertionPointToEnd(&block);
6708	opBuilder.create<vector::YieldOp>(loc, maskedOp->getResults());
6709	}
6710
6711	LogicalResult MaskOp::verify() {
6712	// Structural checks.
6713	Block &block = getMaskRegion().getBlocks().front();
6714	if (block.getOperations().empty())
6715	return emitOpError("expects a terminator within the mask region");
6716
6717	unsigned numMaskRegionOps = block.getOperations().size();
6718	if (numMaskRegionOps > 2)
6719	return emitOpError("expects only one operation to mask");
6720
6721	// Terminator checks.
6722	auto terminator = dyn_cast<vector::YieldOp>(block.back());
6723	if (!terminator)
6724	return emitOpError("expects a terminator within the mask region");
6725
6726	if (terminator->getNumOperands() != getNumResults())
6727	return emitOpError(
6728	"expects number of results to match mask region yielded values");
6729
6730	// Empty vector.mask. Nothing else to check.
6731	if (numMaskRegionOps == 1)
6732	return success();
6733
6734	auto maskableOp = dyn_cast<MaskableOpInterface>(block.front());
6735	if (!maskableOp)
6736	return emitOpError("expects a MaskableOpInterface within the mask region");
6737
6738	// Result checks.
6739	if (maskableOp->getNumResults() != getNumResults())
6740	return emitOpError("expects number of results to match maskable operation "
6741	"number of results");
6742
6743	if (!llvm::equal(maskableOp->getResults(), terminator.getOperands()))
6744	return emitOpError("expects all the results from the MaskableOpInterface "
6745	"to match all the values returned by the terminator");
6746
6747	if (!llvm::equal(maskableOp->getResultTypes(), getResultTypes()))
6748	return emitOpError(
6749	"expects result type to match maskable operation result type");
6750
6751	if (llvm::count_if(maskableOp->getResultTypes(),
6752	[](Type t) { return llvm::isa<VectorType>(t); }) > 1)
6753	return emitOpError("multiple vector results not supported");
6754
6755	// Mask checks.
6756	Type expectedMaskType = maskableOp.getExpectedMaskType();
6757	if (getMask().getType() != expectedMaskType)
6758	return emitOpError("expects a ")
6759	<< expectedMaskType << " mask for the maskable operation";
6760
6761	// Passthru checks.
6762	Value passthru = getPassthru();
6763	if (passthru) {
6764	if (!maskableOp.supportsPassthru())
6765	return emitOpError(
6766	"doesn't expect a passthru argument for this maskable operation");
6767
6768	if (maskableOp->getNumResults() != 1)
6769	return emitOpError("expects result when passthru argument is provided");
6770
6771	if (passthru.getType() != maskableOp->getResultTypes()[0])
6772	return emitOpError("expects passthru type to match result type");
6773	}
6774
6775	return success();
6776	}
6777
6778	/// Folds empty `vector.mask` with no passthru operand and with or without
6779	/// return values. For example:
6780	///
6781	/// %0 = vector.mask %mask { vector.yield %a : vector<8xf32> } :
6782	/// vector<8xi1> -> vector<8xf32>
6783	/// %1 = user_op %0 : vector<8xf32>
6784	///
6785	/// becomes:
6786	///
6787	/// %0 = user_op %a : vector<8xf32>
6788	///
6789	/// Empty `vector.mask` with passthru operand are handled by the canonicalizer
6790	/// as it requires creating new operations.
6791
6792	static LogicalResult foldEmptyMaskOp(MaskOp maskOp, MaskOp::FoldAdaptor adaptor,
6793	SmallVectorImpl<OpFoldResult> &results) {
6794	if (!maskOp.isEmpty() || maskOp.hasPassthru())
6795	return failure();
6796
6797	Block *block = maskOp.getMaskBlock();
6798	auto terminator = cast<vector::YieldOp>(block->front());
6799	if (terminator.getNumOperands() == 0) {
6800	// `vector.mask` has no results, just remove the `vector.mask`.
6801	return success();
6802	}
6803
6804	// `vector.mask` has results, propagate the results.
6805	llvm::append_range(results, terminator.getOperands());
6806	return success();
6807	}
6808
6809	LogicalResult MaskOp::fold(FoldAdaptor adaptor,
6810	SmallVectorImpl<OpFoldResult> &results) {
6811	if (succeeded(foldEmptyMaskOp(*this, adaptor, results)))
6812	return success();
6813
6814	MaskFormat maskFormat = getMaskFormat(getMask());
6815	if (maskFormat != MaskFormat::AllTrue)
6816	return failure();
6817
6818	// Move maskable operation outside of the `vector.mask` region.
6819	Operation *maskableOp = getMaskableOp();
6820	maskableOp->dropAllUses();
6821	maskableOp->moveBefore(getOperation());
6822
6823	llvm::append_range(results, maskableOp->getResults());
6824	return success();
6825	}
6826
6827	/// Canonialize empty `vector.mask` operations that can't be handled in
6828	/// `VectorMask::fold` as they require creating new operations.
6829	///
6830	/// Example 1: Empty `vector.mask` with passthru operand.
6831	///
6832	/// %0 = vector.mask %mask, %passthru { vector.yield %a : vector<8xf32> } :
6833	/// vector<8xi1> -> vector<8xf32>
6834	///
6835	/// becomes:
6836	///
6837	/// %0 = arith.select %mask, %a, %passthru : vector<8xf32>
6838	///
6839	class CanonializeEmptyMaskOp : public OpRewritePattern<MaskOp> {
6840	using OpRewritePattern::OpRewritePattern;
6841
6842	LogicalResult matchAndRewrite(MaskOp maskOp,
6843	PatternRewriter &rewriter) const override {
6844	if (!maskOp.isEmpty())
6845	return failure();
6846
6847	if (!maskOp.hasPassthru())
6848	return failure();
6849
6850	Block *block = maskOp.getMaskBlock();
6851	auto terminator = cast<vector::YieldOp>(block->front());
6852	assert(terminator.getNumOperands() == 1 &&
6853	"expected one result when passthru is provided");
6854
6855	rewriter.replaceOpWithNewOp<arith::SelectOp>(
6856	maskOp, maskOp.getResultTypes(), maskOp.getMask(),
6857	terminator.getOperand(0), maskOp.getPassthru());
6858
6859	return success();
6860	}
6861	};
6862
6863	void MaskOp::getCanonicalizationPatterns(RewritePatternSet &results,
6864	MLIRContext *context) {
6865	results.add<CanonializeEmptyMaskOp>(context);
6866	}
6867
6868	// MaskingOpInterface definitions.
6869
6870	/// Returns the operation masked by this 'vector.mask'.
6871	Operation *MaskOp::getMaskableOp() {
6872	Block *block = getMaskBlock();
6873	if (block->getOperations().size() < 2)
6874	return nullptr;
6875
6876	return &block->front();
6877	}
6878
6879	/// Returns true if 'vector.mask' has a passthru value.
6880	bool MaskOp::hasPassthru() { return getPassthru() != Value(); }
6881
6882	//===----------------------------------------------------------------------===//
6883	// ScanOp
6884	//===----------------------------------------------------------------------===//
6885
6886	LogicalResult ScanOp::verify() {
6887	VectorType srcType = getSourceType();
6888	VectorType initialType = getInitialValueType();
6889	// Check reduction dimension < rank.
6890	int64_t srcRank = srcType.getRank();
6891	int64_t reductionDim = getReductionDim();
6892	if (reductionDim >= srcRank)
6893	return emitOpError("reduction dimension ")
6894	<< reductionDim << " has to be less than " << srcRank;
6895
6896	// Check that rank(initial_value) = rank(src) - 1.
6897	int64_t initialValueRank = initialType.getRank();
6898	if (initialValueRank != srcRank - 1)
6899	return emitOpError("initial value rank ")
6900	<< initialValueRank << " has to be equal to " << srcRank - 1;
6901
6902	// Check shapes of initial value and src.
6903	ArrayRef<int64_t> srcShape = srcType.getShape();
6904	ArrayRef<int64_t> initialValueShapes = initialType.getShape();
6905	SmallVector<int64_t> expectedShape;
6906	for (int i = 0; i < srcRank; i++) {
6907	if (i != reductionDim)
6908	expectedShape.push_back(srcShape[i]);
6909	}
6910	if (!llvm::equal(initialValueShapes, expectedShape)) {
6911	return emitOpError("incompatible input/initial value shapes");
6912	}
6913
6914	// Verify supported reduction kind.
6915	Type eltType = getDestType().getElementType();
6916	if (!isSupportedCombiningKind(getKind(), eltType))
6917	return emitOpError("unsupported reduction type ")
6918	<< eltType << " for kind '" << stringifyCombiningKind(getKind())
6919	<< "'";
6920
6921	return success();
6922	}
6923
6924	void mlir::vector::populateVectorToVectorCanonicalizationPatterns(
6925	RewritePatternSet &patterns, PatternBenefit benefit) {
6926	patterns
6927	.add<CreateMaskFolder, MaskedLoadFolder, MaskedStoreFolder, GatherFolder,
6928	ScatterFolder, ExpandLoadFolder, CompressStoreFolder,
6929	StridedSliceConstantMaskFolder, TransposeFolder>(
6930	arg: patterns.getContext(), args&: benefit);
6931	}
6932
6933	//===----------------------------------------------------------------------===//
6934	// SplatOp
6935	//===----------------------------------------------------------------------===//
6936
6937	OpFoldResult SplatOp::fold(FoldAdaptor adaptor) {
6938	auto constOperand = adaptor.getInput();
6939	if (!isa_and_nonnull<IntegerAttr, FloatAttr>(constOperand))
6940	return {};
6941
6942	// SplatElementsAttr::get treats single value for second arg as being a splat.
6943	return SplatElementsAttr::get(getType(), {constOperand});
6944	}
6945
6946	void SplatOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
6947	SetIntRangeFn setResultRanges) {
6948	setResultRanges(getResult(), argRanges.front());
6949	}
6950
6951	Value mlir::vector::makeArithReduction(OpBuilder &b, Location loc,
6952	CombiningKind kind, Value v1, Value acc,
6953	arith::FastMathFlagsAttr fastmath,
6954	Value mask) {
6955	Type t1 = getElementTypeOrSelf(type: v1.getType());
6956	Type tAcc = getElementTypeOrSelf(type: acc.getType());
6957	Value result;
6958
6959	switch (kind) {
6960	case CombiningKind::ADD:
6961	if (t1.isIntOrIndex() && tAcc.isIntOrIndex())
6962	result = b.createOrFold<arith::AddIOp>(loc, v1, acc);
6963	else if (llvm::isa<FloatType>(Val: t1) && llvm::isa<FloatType>(Val: tAcc))
6964	result = b.createOrFold<arith::AddFOp>(loc, v1, acc, fastmath);
6965	else
6966	llvm_unreachable("invalid value types for ADD reduction");
6967	break;
6968	case CombiningKind::AND:
6969	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
6970	result = b.createOrFold<arith::AndIOp>(loc, v1, acc);
6971	break;
6972	case CombiningKind::MAXNUMF:
6973	assert(llvm::isa<FloatType>(t1) && llvm::isa<FloatType>(tAcc) &&
6974	"expected float values");
6975	result = b.createOrFold<arith::MaxNumFOp>(loc, v1, acc, fastmath);
6976	break;
6977	case CombiningKind::MAXIMUMF:
6978	assert(llvm::isa<FloatType>(t1) && llvm::isa<FloatType>(tAcc) &&
6979	"expected float values");
6980	result = b.createOrFold<arith::MaximumFOp>(loc, v1, acc, fastmath);
6981	break;
6982	case CombiningKind::MINNUMF:
6983	assert(llvm::isa<FloatType>(t1) && llvm::isa<FloatType>(tAcc) &&
6984	"expected float values");
6985	result = b.createOrFold<arith::MinNumFOp>(loc, v1, acc, fastmath);
6986	break;
6987	case CombiningKind::MINIMUMF:
6988	assert(llvm::isa<FloatType>(t1) && llvm::isa<FloatType>(tAcc) &&
6989	"expected float values");
6990	result = b.createOrFold<arith::MinimumFOp>(loc, v1, acc, fastmath);
6991	break;
6992	case CombiningKind::MAXSI:
6993	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
6994	result = b.createOrFold<arith::MaxSIOp>(loc, v1, acc);
6995	break;
6996	case CombiningKind::MINSI:
6997	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
6998	result = b.createOrFold<arith::MinSIOp>(loc, v1, acc);
6999	break;
7000	case CombiningKind::MAXUI:
7001	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
7002	result = b.createOrFold<arith::MaxUIOp>(loc, v1, acc);
7003	break;
7004	case CombiningKind::MINUI:
7005	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
7006	result = b.createOrFold<arith::MinUIOp>(loc, v1, acc);
7007	break;
7008	case CombiningKind::MUL:
7009	if (t1.isIntOrIndex() && tAcc.isIntOrIndex())
7010	result = b.createOrFold<arith::MulIOp>(loc, v1, acc);
7011	else if (llvm::isa<FloatType>(Val: t1) && llvm::isa<FloatType>(Val: tAcc))
7012	result = b.createOrFold<arith::MulFOp>(loc, v1, acc, fastmath);
7013	else
7014	llvm_unreachable("invalid value types for MUL reduction");
7015	break;
7016	case CombiningKind::OR:
7017	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
7018	result = b.createOrFold<arith::OrIOp>(loc, v1, acc);
7019	break;
7020	case CombiningKind::XOR:
7021	assert(t1.isIntOrIndex() && tAcc.isIntOrIndex() && "expected int values");
7022	result = b.createOrFold<arith::XOrIOp>(loc, v1, acc);
7023	break;
7024	};
7025
7026	assert(result && "unknown CombiningKind");
7027	return selectPassthru(builder&: b, mask, newValue: result, passthru: acc);
7028	}
7029
7030	//===----------------------------------------------------------------------===//
7031	// Vector Masking Utilities
7032	//===----------------------------------------------------------------------===//
7033
7034	/// Create the vector.yield-ended region of a vector.mask op with `maskableOp`
7035	/// as masked operation.
7036	void mlir::vector::createMaskOpRegion(OpBuilder &builder,
7037	Operation *maskableOp) {
7038	assert(maskableOp->getBlock() && "MaskableOp must be inserted into a block");
7039	Block *insBlock = builder.getInsertionBlock();
7040	// Create a block and move the op to that block.
7041	insBlock->getOperations().splice(
7042	where: insBlock->begin(), L2&: maskableOp->getBlock()->getOperations(), N: maskableOp);
7043	builder.create<YieldOp>(maskableOp->getLoc(), maskableOp->getResults());
7044	}
7045
7046	/// Creates a vector.mask operation around a maskable operation. Returns the
7047	/// vector.mask operation if the mask provided is valid. Otherwise, returns
7048	/// the maskable operation itself.
7049	Operation *mlir::vector::maskOperation(OpBuilder &builder,
7050	Operation *maskableOp, Value mask,
7051	Value passthru) {
7052	if (!mask)
7053	return maskableOp;
7054	if (passthru)
7055	return builder.create<MaskOp>(maskableOp->getLoc(),
7056	maskableOp->getResultTypes(), mask, passthru,
7057	maskableOp, createMaskOpRegion);
7058	return builder.create<MaskOp>(maskableOp->getLoc(),
7059	maskableOp->getResultTypes(), mask, maskableOp,
7060	createMaskOpRegion);
7061	}
7062
7063	/// Creates a vector select operation that picks values from `newValue` or
7064	/// `passthru` for each result vector lane based on `mask`. This utility is used
7065	/// to propagate the pass-thru value of vector.mask or for cases where only the
7066	/// pass-thru value propagation is needed. VP intrinsics do not support
7067	/// pass-thru values and every mask-out lane is set to poison. LLVM backends are
7068	/// usually able to match op + select patterns and fold them into a native
7069	/// target instructions.
7070	Value mlir::vector::selectPassthru(OpBuilder &builder, Value mask,
7071	Value newValue, Value passthru) {
7072	if (!mask)
7073	return newValue;
7074
7075	return builder.create<arith::SelectOp>(newValue.getLoc(), newValue.getType(),
7076	mask, newValue, passthru);
7077	}
7078
7079	//===----------------------------------------------------------------------===//
7080	// TableGen'd op method definitions
7081	//===----------------------------------------------------------------------===//
7082
7083	#define GET_ATTRDEF_CLASSES
7084	#include "mlir/Dialect/Vector/IR/VectorAttributes.cpp.inc"
7085
7086	#define GET_OP_CLASSES
7087	#include "mlir/Dialect/Vector/IR/VectorOps.cpp.inc"
7088