[Mlir-commits] [mlir] [MLIR][XeGPU] Add lowering from transfer_read/transfer_write to load_gather/store_scatter (PR #152429)
Jianhui Li
llvmlistbot at llvm.org
Wed Aug 13 14:50:09 PDT 2025
================
@@ -155,6 +152,300 @@ createNdDescriptor(PatternRewriter &rewriter, Location loc,
return ndDesc;
}
+// Adjusts the strides of a memref according to a given permutation map for
+// vector operations.
+//
+// This function updates the last `vecRank` elements of the `strides` array to
+// reflect the permutation specified by `permMap`. The permutation is applied
+// to the innermost dimensions of the memref, corresponding to the vector
+// shape. This is typically used when lowering vector transfer operations with
+// permutation maps to memory accesses, ensuring that the memory strides match
+// the logical permutation of vector dimensions.
+//
+// Example:
+// Suppose we have a memref of rank 4 with strides `[s0, s1, s2, s3]` and a
+// vector of rank 2. If the permutation map swaps the last two dimensions
+// (e.g., [0, 1] -> [1, 0]), then after calling this function, the last two
+// strides will be swapped:
+// Original strides: [s0, s1, s2, s3]
+// After permutation: [s0, s1, s3, s2]
+//
+void adjustStridesForPermutation(Operation *op, PatternRewriter &rewriter,
+ MemRefType memrefType, AffineMap permMap,
+ VectorType vecType,
+ SmallVectorImpl<Value> &strides) {
+ unsigned vecRank;
+ unsigned memrefRank = memrefType.getRank();
+
+ if (permMap.isMinorIdentity())
+ return;
+ vecRank = vecType.getRank();
+ // Only adjust the last vecRank strides according to the permutation
+ ArrayRef<Value> relevantStrides = ArrayRef<Value>(strides).take_back(vecRank);
+ SmallVector<Value> adjustedStrides(vecRank);
+ // For each output dimension in the permutation map, find which input dim it
+ // refers to, and assign the corresponding stride.
+ for (unsigned outIdx = 0; outIdx < vecRank; ++outIdx) {
+ AffineExpr expr = permMap.getResult(outIdx);
+ auto dimExpr = dyn_cast<AffineDimExpr>(expr);
+ assert(dimExpr && "The permutation expr must be affine expression");
+ unsigned pos = dimExpr.getPosition();
+ // Map permutation to the relevant strides (innermost dims)
+ assert((pos >= (memrefRank - vecRank)) &&
+ "Permuted index must be in the inner dimensions");
+
+ // The stride for output dimension outIdx is the stride of input dimension
+ // pos
+ adjustedStrides[outIdx] = relevantStrides[pos - (memrefRank - vecRank)];
+ }
+ // Replace the last vecRank strides with the adjusted ones
+ for (unsigned i = 0; i < vecRank; ++i)
+ strides[memrefRank - vecRank + i] = adjustedStrides[i];
+}
+
+SmallVector<Value> computeStrides(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter) {
+ SmallVector<Value> strides;
+ Value baseMemref = xferOp.getBase();
+ AffineMap permMap = xferOp.getPermutationMap();
+ VectorType vectorType = xferOp.getVectorType();
+ MemRefType memrefType = dyn_cast<MemRefType>(baseMemref.getType());
+
+ Location loc = xferOp.getLoc();
+ if (memrefType.hasStaticShape()) {
+ int64_t offset;
+ SmallVector<int64_t> intStrides;
+ if (failed(memrefType.getStridesAndOffset(intStrides, offset)))
+ return {};
+ // Wrap static strides as MLIR values
+ for (int64_t s : intStrides)
+ strides.push_back(arith::ConstantIndexOp::create(rewriter, loc, s));
+ } else {
+ // For dynamic shape memref, use memref.extract_strided_metadata to get
+ // stride values
+ unsigned rank = memrefType.getRank();
+ Type indexType = rewriter.getIndexType();
+
+ // Result types: [base_memref, offset, stride0, stride1, ..., strideN-1,
+ // size0, size1, ..., sizeN-1]
+ SmallVector<Type> resultTypes;
+ resultTypes.push_back(MemRefType::get(
+ {}, memrefType.getElementType())); // base memref (unranked)
+ resultTypes.push_back(indexType); // offset
+ for (unsigned i = 0; i < rank; ++i) {
+ resultTypes.push_back(indexType); // strides
+ }
+ for (unsigned i = 0; i < rank; ++i) {
+ resultTypes.push_back(indexType); // sizes
+ }
+
+ auto meta = memref::ExtractStridedMetadataOp::create(
+ rewriter, loc, resultTypes, baseMemref);
+ strides.append(meta.getStrides().begin(), meta.getStrides().end());
+ }
+ // Adjust strides according to the permutation map (e.g., for transpose)
+ adjustStridesForPermutation(xferOp, rewriter, memrefType, permMap, vectorType,
+ strides);
+ return strides;
+}
+
+// This function compute the vectors of localOffsets for scattered load/stores.
+// It is used in the lowering of vector.transfer_read/write to
+// load_gather/store_scatter Example:
+// %0 = vector.transfer_read %expand_shape[%block_id_y, %c0, %c0, %c0, %c0],
+// %cst {in_bounds = [true, true, true, true]}>} :
+// memref<8x4x2x6x32xbf16>, vector<4x2x6x32xbf16>
+//
+// %6 = vector.step: vector<4xindex>
+// %7 = vector.step: vector<2xindex>
+// %8 = vector.step: vector<6xindex>
+// %9 = vector.step: vector<32xindex>
+// %10 = arith.mul %6, 384
+// %11 = arith.mul %7, 192
+// %12 = arith.mul %8, 32
+// %13 = arith.mul %9, 1
+// %14 = vector.shape_cast %10: vector<4xindex> -> vector<4x1x1x1xbf16>
+// %15 = vector.shape_cast %11: vector<2xindex> -> vector<1x2x1x1xbf16>
+// %16 = vector.shape_cast %12: vector<6xindex> -> vector<1x1x6x1xbf16>
+// %17 = vector.shape_cast %13: vector<32xindex> -> vector<1x1x1x32xbf16>
+// %18 = vector.broadcast %14: vector<4x1x1x1xbf16> -> vector<4x2x6x32xindex>
+// %19 = vector.broadcast %15: vector<1x2x1x1xbf16> -> vector<4x2x6x32xindex>
+// %20 = vector.broadcast %16: vector<1x1x6x1xbf16> -> vector<4x2x6x32xindex>
+// %21 = vector.broadcast %17: vector<1x1x1x32xbf16> -> vector<4x2x6x32xindex>
+// %22 = arith.add %18, %19
+// %23 = arith.add %20, %21
+// %local_offsets = arith.add %22, %23
+// %orig_offset = %block_id_y * 4x2x6x32 // consider using affine map
+// %offsets = orig_offset + local_offsets
+static Value computeOffsets(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter,
+ ArrayRef<Value> strides) {
+ Location loc = xferOp.getLoc();
+ VectorType vectorType = xferOp.getVectorType();
+ SmallVector<Value> indices(xferOp.getIndices().begin(),
+ xferOp.getIndices().end());
+ ArrayRef<int64_t> vectorShape = vectorType.getShape();
+
+ // Create vector.step operations for each dimension
+ SmallVector<Value> stepVectors;
+ for (int64_t dim : vectorShape) {
+ auto stepType = VectorType::get({dim}, rewriter.getIndexType());
+ auto stepOp = vector::StepOp::create(rewriter, loc, stepType);
+ stepVectors.push_back(stepOp);
+ }
+
+ // Multiply step vectors by corresponding strides
+ size_t memrefRank = strides.size();
+ size_t vectorRank = vectorShape.size();
+ SmallVector<Value> strideMultiplied;
+ for (size_t i = 0; i < vectorRank; ++i) {
+ size_t memrefDim = memrefRank - vectorRank + i;
+ Value strideValue = strides[memrefDim];
+ auto mulType = dyn_cast<VectorType>(stepVectors[i].getType());
+ auto bcastOp =
+ vector::BroadcastOp::create(rewriter, loc, mulType, strideValue);
+ auto mulOp = arith::MulIOp::create(rewriter, loc, stepVectors[i], bcastOp);
+ strideMultiplied.push_back(mulOp);
+ }
+
+ // Shape cast each multiplied vector to add singleton dimensions
+ SmallVector<Value> shapeCasted;
+ for (size_t i = 0; i < vectorRank; ++i) {
+ SmallVector<int64_t> newShape(vectorRank, 1);
+ newShape[i] = vectorShape[i];
+ auto newType = VectorType::get(newShape, rewriter.getIndexType());
+ auto castOp = vector::ShapeCastOp::create(rewriter, loc, newType,
+ strideMultiplied[i]);
+ shapeCasted.push_back(castOp);
+ }
+
+ // Broadcast each shape-casted vector to full vector shape
+ SmallVector<Value> broadcasted;
+ auto fullIndexVectorType =
+ VectorType::get(vectorShape, rewriter.getIndexType());
+ for (Value shapeCastVal : shapeCasted) {
+ auto broadcastOp = vector::BroadcastOp::create(
+ rewriter, loc, fullIndexVectorType, shapeCastVal);
+ broadcasted.push_back(broadcastOp);
+ }
+
+ // Add all broadcasted vectors together to compute local offsets
+ Value localOffsets = broadcasted[0];
+ for (size_t i = 1; i < broadcasted.size(); ++i) {
+ localOffsets =
+ arith::AddIOp::create(rewriter, loc, localOffsets, broadcasted[i]);
+ }
+
+ // Compute base offset from transfer read indices
+ Value baseOffset = nullptr;
+ if (!indices.empty()) {
+ baseOffset = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ for (size_t i = 0; i < indices.size(); ++i) {
+ Value strideVal = strides[i];
+ Value offsetContrib =
+ arith::MulIOp::create(rewriter, loc, indices[i], strideVal);
+ baseOffset =
+ arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
+ }
+ // Broadcast base offset to match vector shape
+ Value bcastBase = vector::BroadcastOp::create(
+ rewriter, loc, fullIndexVectorType, baseOffset);
+ localOffsets =
+ arith::AddIOp::create(rewriter, loc, bcastBase, localOffsets);
+ }
+ return localOffsets;
+}
+
+// Collapse memref shape to 1D
+static Value collapseMemrefTo1D(VectorTransferOpInterface xferOp,
+ PatternRewriter &rewriter) {
+ Location loc = xferOp.getLoc();
+
+ Value baseMemref = xferOp.getBase();
+ MemRefType memrefType = dyn_cast<MemRefType>(baseMemref.getType());
+ Type elementType = memrefType.getElementType();
+
+ // Compute the total number of elements in the memref
+ MemRefType flatMemrefType;
+ if (memrefType.hasStaticShape()) {
+ auto totalElements = memrefType.getNumElements();
+ flatMemrefType = MemRefType::get({totalElements}, elementType);
+ } else {
+ flatMemrefType = MemRefType::get({ShapedType::kDynamic}, elementType);
+ }
+
+ SmallVector<ReassociationIndices> reassociation;
+ ReassociationIndices allDims =
+ llvm::to_vector(llvm::seq<int64_t>(0, memrefType.getRank()));
+ reassociation.push_back(allDims);
+
+ auto collapseOp = memref::CollapseShapeOp::create(
+ rewriter, loc, flatMemrefType, baseMemref, reassociation);
+ return collapseOp;
+}
+
+LogicalResult lowerToRegularLoadOp(vector::TransferReadOp readOp,
----------------
Jianhui-Li wrote:
added
https://github.com/llvm/llvm-project/pull/152429
More information about the Mlir-commits
mailing list