| //===- XeGPUPeepHoleOptimizer.cpp - XeGPU optimize block loads -*- C++ -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
| #include "mlir/Dialect/MemRef/IR/MemRef.h" |
| #include "mlir/Dialect/SCF/Transforms/Patterns.h" |
| #include "mlir/Dialect/Utils/IndexingUtils.h" |
| #include "mlir/Dialect/Utils/StaticValueUtils.h" |
| #include "mlir/Dialect/Vector/IR/VectorOps.h" |
| #include "mlir/Dialect/XeGPU/IR/XeGPU.h" |
| #include "mlir/Dialect/XeGPU/Transforms/Passes.h" |
| #include "mlir/Dialect/XeGPU/Transforms/Transforms.h" |
| #include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h" |
| #include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h" |
| #include "mlir/Dialect/XeGPU/uArch/uArchBase.h" |
| #include "mlir/Dialect/XeGPU/uArch/uArchCommon.h" |
| #include "mlir/IR/BuiltinTypes.h" |
| #include "mlir/IR/OpDefinition.h" |
| #include "mlir/IR/Types.h" |
| #include "mlir/IR/Value.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "llvm/ADT/SmallVector.h" |
| #include <optional> |
| |
| namespace mlir { |
| namespace xegpu { |
| #define GEN_PASS_DEF_XEGPUPEEPHOLEOPTIMIZER |
| #include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc" |
| } // namespace xegpu |
| } // namespace mlir |
| |
| #define DEBUG_TYPE "xegpu-optimize-peephole" |
| #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") |
| |
| using namespace mlir; |
| |
| namespace { |
| |
| /// Get the 2D lane data from a tensor desc type if it exists. |
| static std::optional<SmallVector<int64_t>> |
| getMaybeLaneData(xegpu::TensorDescType tdescType) { |
| auto layout = tdescType.getLayoutAttr(); |
| if (!layout) |
| return std::nullopt; |
| return xegpu::getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt()); |
| } |
| |
| /// Get the 2D lane layout from a tensor desc type if it exists. |
| static std::optional<SmallVector<int64_t>> |
| getMaybeLaneLayout(xegpu::TensorDescType tdescType) { |
| auto layout = tdescType.getLayoutAttr(); |
| if (!layout) |
| return std::nullopt; |
| return xegpu::getInner2DIfUnitLeadingDims( |
| layout.getEffectiveLaneLayoutAsInt()); |
| } |
| |
| /// A layout can be optimized if its lane layout is transposed (lane[0] != 1 && |
| /// lane[1] == 1), but inner lane data is not equal to [1, 1]. |
| /// Example: |
| /// !xegpu.tensor_desc<16x16xf16, |
| /// #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>> |
| /// In this case, lane layout is transposed (from the usual [1, SG_SIZE] form) |
| /// indicating that this is a load that requires transpose effect. However, |
| /// lane data is [1, 2], meaning that each lane must grab 2 f16 elements from |
| /// the inner dimension. We convert this to a optimized form by converting the |
| /// tensor_desc to i32 type such that lane data becomes [1, 1]. This makes the |
| /// later lowering easily use the load with transpose instruction. |
| static bool canBeOptimizedForTranspose(ArrayRef<int64_t> laneLayout, |
| ArrayRef<int64_t> laneData) { |
| if (laneLayout.size() != 2 || laneData.size() != 2) |
| return false; |
| if (laneLayout[0] == 1 || laneLayout[1] != 1) |
| return false; |
| if (laneData[0] != 1 || laneData[1] == 1) |
| return false; |
| return true; |
| } |
| |
| /// A tensor desc type can be optimized if its element type is less than 32 bits |
| /// and its layout can be optimized. |
| static bool canBeOptimizedForTranspose(xegpu::TensorDescType tdescType) { |
| // If the dtype is greater or equal to 32 bits, layout must be valid. |
| int elementTyBitwidth = tdescType.getElementType().getIntOrFloatBitWidth(); |
| if (elementTyBitwidth >= 32) |
| return false; |
| auto maybeLaneLayout = getMaybeLaneLayout(tdescType); |
| auto maybeLaneData = getMaybeLaneData(tdescType); |
| if (!maybeLaneData || !maybeLaneLayout) |
| return false; |
| return canBeOptimizedForTranspose(*maybeLaneLayout, *maybeLaneData); |
| } |
| |
| /// Check if a tensor desc type can be optimized for transpose, if so return the |
| /// new optimized tensor desc type with a valid transpose layout. |
| static xegpu::TensorDescType |
| tryOptimize(xegpu::TensorDescType tdescType, |
| const xegpu::uArch::uArch *targetuArch) { |
| if (!canBeOptimizedForTranspose(tdescType)) |
| return tdescType; |
| auto laneData = getMaybeLaneData(tdescType) |
| .value(); // Lane data must exist if we reach here. |
| int64_t innerLaneData = laneData[1]; |
| int elementTyBitwidth = tdescType.getElementType().getIntOrFloatBitWidth(); |
| // Required shape is total shape of the vector result that this tensor desc |
| // must eventually load after adjusting for the new bitwidth and array |
| // length. |
| SmallVector<int64_t> requiredShape(tdescType.getShape()); |
| requiredShape.back() = |
| requiredShape.back() * tdescType.getArrayLength() / innerLaneData; |
| int newBitWidth = elementTyBitwidth * innerLaneData; |
| Type newElemTy = IntegerType::get(tdescType.getContext(), newBitWidth); |
| // Supported shape is the max transpose shape that can be supported by |
| // hardware that is less than or equal to required shape. |
| auto *blockLoadTarget = |
| dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>( |
| targetuArch->getInstruction( |
| xegpu::uArch::InstructionKind::Subgroup2DBlockLoad)); |
| auto maybeHWParams = blockLoadTarget->getBlockWidthHeightCount( |
| newElemTy, /** has transform */ false, /** has transpose */ true); |
| // If no HW params found, return the original type. |
| if (!maybeHWParams) |
| return tdescType; |
| auto [widths, heights, counts] = maybeHWParams.value(); |
| // TODO: Currently we expect array length to be 1 for transpose case. |
| if (counts.size() != 1 || counts[0] != 1) |
| return tdescType; |
| int arrayLen = counts[0]; |
| // The transpose acts on the innermost 2 dims; any leading dims are unit |
| // batch. |
| int64_t rank = requiredShape.size(); |
| int supportedHeight = xegpu::getLargestDivisor( |
| static_cast<int>(requiredShape[rank - 2]), heights); |
| int supportedWidth = xegpu::getLargestDivisor( |
| static_cast<int>(requiredShape[rank - 1]), widths); |
| // If no supported height or width found, return the original type. |
| if (supportedHeight == -1 || supportedWidth == -1) |
| return tdescType; |
| |
| // Preserve leading (unit) batch dims; only the inner 2 dims are reshaped. |
| SmallVector<int64_t> supportedShape(requiredShape.begin(), |
| requiredShape.end() - 2); |
| supportedShape.push_back(supportedHeight); |
| supportedShape.push_back(supportedWidth); |
| auto ctx = tdescType.getContext(); |
| auto origLayout = tdescType.getLayoutAttr(); |
| auto laneLayoutI64 = origLayout.getEffectiveLaneLayoutAsInt(); |
| SmallVector<int32_t> laneLayoutI32(laneLayoutI64.begin(), |
| laneLayoutI64.end()); |
| |
| xegpu::LayoutAttr newLayout = xegpu::LayoutAttr::get( |
| ctx, /*lane_layout=*/DenseI32ArrayAttr::get(ctx, laneLayoutI32), |
| /*lane_data=*/ |
| DenseI32ArrayAttr::get(ctx, |
| SmallVector<int32_t>(laneLayoutI32.size(), 1)), |
| /*order=*/origLayout.getOrder()); |
| |
| // Array length can not be larger than 1 for transpose case. |
| return xegpu::TensorDescType::get(supportedShape, newElemTy, arrayLen, |
| tdescType.getBoundaryCheck(), |
| tdescType.getMemorySpace(), newLayout); |
| } |
| |
| /// Helper to convert an OpFoldResult to Value. |
| static Value convertToValue(ConversionPatternRewriter &rewriter, Location loc, |
| OpFoldResult ofr) { |
| std::optional<int64_t> mayBeInt = getConstantIntValue(ofr); |
| if (mayBeInt) |
| return arith::ConstantIndexOp::create(rewriter, loc, *mayBeInt).getResult(); |
| return llvm::cast<Value>(ofr); |
| } |
| |
| /// Helper to divide a Value by a constant integer. |
| static Value divideByConstant(ConversionPatternRewriter &rewriter, Location loc, |
| Value val, int64_t constant) { |
| // If the constant is a power of 2, use right shift for division. |
| if (llvm::isPowerOf2_64(constant)) { |
| int64_t shiftAmount = llvm::Log2_64(constant); |
| return arith::ShRUIOp::create( |
| rewriter, loc, val, |
| arith::ConstantIndexOp::create(rewriter, loc, shiftAmount) |
| .getResult()) |
| .getResult(); |
| } |
| auto constantOp = |
| arith::ConstantIndexOp::create(rewriter, loc, constant).getResult(); |
| return arith::DivUIOp::create(rewriter, loc, val, constantOp).getResult(); |
| } |
| |
| /// This function takes a larger register block `data` and generates multiple |
| /// smaller loads (size given by `newTensorDesc`) to fill in the `data` block |
| /// starting from `offsets`. |
| static Value generateLoads(ConversionPatternRewriter &rewriter, |
| TypedValue<VectorType> data, |
| SmallVector<OpFoldResult> offsets, |
| TypedValue<xegpu::TensorDescType> newTensorDesc, |
| xegpu::LoadNdOp origLoadOp) { |
| Location loc = data.getLoc(); |
| assert(offsets.size() >= 2 && "Expecting at least 2 offsets for 2D LoadNdOp"); |
| int64_t rank = data.getType().getRank(); |
| Value offsetDim0 = convertToValue(rewriter, loc, offsets[offsets.size() - 2]); |
| Value offsetDim1 = convertToValue(rewriter, loc, offsets[offsets.size() - 1]); |
| SmallVector<int64_t> supportedShape(newTensorDesc.getType().getShape()); |
| // Compute the ratio between original shape and supported shape. We need to |
| // generate loads in this ratio arrangement. |
| auto shapeRatio = computeShapeRatio(data.getType().getShape(), |
| supportedShape) |
| .value(); // `ratio` must be defined if we reach here. |
| // The loop below only walks the last 2 entries of `shapeRatio`; leading |
| // (batch) dims are unit, so their ratio is 1 and their offsets pass through |
| // unchanged. |
| int64_t suppDim0 = supportedShape[rank - 2]; |
| int64_t suppDim1 = supportedShape[rank - 1]; |
| for (int64_t h = 0; h < shapeRatio[rank - 2]; ++h) { |
| for (int64_t w = 0; w < shapeRatio[rank - 1]; ++w) { |
| int64_t localOffsetDim0 = h * suppDim0; |
| int64_t localOffsetDim1 = w * suppDim1; |
| Value loadOffsetX = arith::AddIOp::create( |
| rewriter, loc, offsetDim0, |
| arith::ConstantIndexOp::create(rewriter, loc, localOffsetDim0) |
| .getResult()); |
| Value loadOffsetY = arith::AddIOp::create( |
| rewriter, loc, offsetDim1, |
| arith::ConstantIndexOp::create(rewriter, loc, localOffsetDim1) |
| .getResult()); |
| // Keep the leading (batch) offsets; replace only the inner 2. |
| SmallVector<OpFoldResult> loadOffsets(offsets.begin(), offsets.end()); |
| loadOffsets[loadOffsets.size() - 2] = loadOffsetX; |
| loadOffsets[loadOffsets.size() - 1] = loadOffsetY; |
| auto loadOp = xegpu::LoadNdOp::create( |
| rewriter, loc, |
| VectorType::get(supportedShape, data.getType().getElementType()), |
| newTensorDesc, loadOffsets, origLoadOp.getPackedAttr(), |
| origLoadOp.getTransposeAttr(), origLoadOp.getL1HintAttr(), |
| origLoadOp.getL2HintAttr(), origLoadOp.getL3HintAttr(), |
| origLoadOp.getLayoutAttr()); |
| // Set the layout for the loadOp. |
| auto layoutAttr = newTensorDesc.getType().getLayoutAttr(); |
| loadOp.setAnchorLayout(layoutAttr); |
| // Insert the loaded block into the right position in data (leading dims |
| // at 0, inner 2 dims at the local tile offset). |
| SmallVector<int64_t> insertPos(rank, 0); |
| insertPos[rank - 2] = localOffsetDim0; |
| insertPos[rank - 1] = localOffsetDim1; |
| SmallVector<int64_t> insertStrides(rank, 1); |
| auto insertOp = vector::InsertStridedSliceOp::create( |
| rewriter, loc, loadOp.getResult(), data, insertPos, insertStrides); |
| // InsertOp must have the same layout as newTensorDesc. |
| xegpu::setTemporaryLayout(insertOp->getOpResult(0), layoutAttr); |
| data = insertOp.getResult(); |
| } |
| } |
| return data; |
| } |
| |
| /// Checks if a CreateNdDescOp can be optimized for transpose, if so creates a |
| /// new CreateNdDescOp with optimized tensor desc type. This involves extracting |
| /// the base pointer from the original memory source and adjusting the shape and |
| /// strides of the tensor desc to fit with the new optimized transpose layout. |
| class XeGPUCreateNdDescOpPattern final |
| : public OpConversionPattern<xegpu::CreateNdDescOp> { |
| public: |
| using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern; |
| LogicalResult |
| matchAndRewrite(xegpu::CreateNdDescOp createNdOp, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto tdescTy = createNdOp.getType(); |
| // Get the target uArch info. |
| auto chipStr = xegpu::getChipStr(createNdOp); |
| // Check if the chip is supported. |
| assert(chipStr && |
| (chipStr.value() == "pvc" || chipStr.value() == "bmg" || |
| chipStr.value() == "cri") && |
| "Expecting target chip to be pvc, bmg or cri for transpose " |
| "optimization."); |
| const auto *targetuArch = xegpu::uArch::getUArch(chipStr.value()); |
| |
| auto convertType = tryOptimize(tdescTy, targetuArch); |
| if (convertType == tdescTy) |
| return failure(); |
| Location loc = createNdOp.getLoc(); |
| Value source = createNdOp.getSource(); |
| auto memrefType = dyn_cast<MemRefType>(source.getType()); |
| |
| bool dynamicMemref = |
| memrefType && !xegpu::hasStaticShapeAndStrides(memrefType); |
| SmallVector<OpFoldResult> mixedSizes; |
| SmallVector<OpFoldResult> mixedStrides; |
| memref::ExtractStridedMetadataOp meta; |
| if (dynamicMemref) { |
| meta = memref::ExtractStridedMetadataOp::create(rewriter, loc, source); |
| mixedSizes = meta.getConstifiedMixedSizes(); |
| mixedStrides = meta.getConstifiedMixedStrides(); |
| } else { |
| mixedSizes = createNdOp.getMixedSizes(); |
| mixedStrides = createNdOp.getMixedStrides(); |
| } |
| |
| auto maybeConstInnerStride = getConstantIntValue(mixedStrides.back()); |
| // Only row-major memrefs are expected for now. |
| if (!maybeConstInnerStride || *maybeConstInnerStride != 1) |
| return rewriter.notifyMatchFailure( |
| createNdOp, "Expecting row-major memref for transpose optimization."); |
| auto optionalLaneData = getMaybeLaneData(tdescTy); |
| assert(optionalLaneData && "Expected 2D lane data"); |
| auto laneData = optionalLaneData.value(); |
| int64_t innerLaneData = laneData[1]; |
| // Inner dimension of the shape must be adjusted based on innerLaneData. |
| SmallVector<OpFoldResult> modifiedShape(mixedSizes); |
| modifiedShape.back() = divideByConstant( |
| rewriter, loc, convertToValue(rewriter, loc, modifiedShape.back()), |
| innerLaneData); |
| // Repacking to a wider element rescales every stride but the innermost. |
| assert(mixedStrides.size() >= 2 && |
| "Expected at least 2 strides for CreateNdDescOp"); |
| SmallVector<OpFoldResult> modifiedStrides(mixedStrides); |
| for (size_t i = 0; i + 1 < modifiedStrides.size(); ++i) |
| modifiedStrides[i] = divideByConstant( |
| rewriter, loc, convertToValue(rewriter, loc, modifiedStrides[i]), |
| innerLaneData); |
| |
| if (memrefType) { |
| Value baseIdx; |
| if (dynamicMemref) { |
| // Base = aligned base pointer + structural offset (in bytes). |
| Value alignedPtr = memref::ExtractAlignedPointerAsIndexOp::create( |
| rewriter, loc, meta.getBaseBuffer()); |
| Value elemBytes = arith::ConstantIndexOp::create( |
| rewriter, loc, memrefType.getElementTypeBitWidth() / 8); |
| Value offBytes = |
| arith::MulIOp::create(rewriter, loc, meta.getOffset(), elemBytes); |
| baseIdx = arith::AddIOp::create(rewriter, loc, alignedPtr, offBytes); |
| } else { |
| baseIdx = memref::ExtractAlignedPointerAsIndexOp::create(rewriter, loc, |
| source); |
| } |
| source = arith::IndexCastOp::create(rewriter, loc, rewriter.getI64Type(), |
| baseIdx); |
| } |
| // Create a new CreateNdDescOp with the modified shape and converted type. |
| auto newCreateNdDescOp = xegpu::CreateNdDescOp::create( |
| rewriter, loc, convertType, source, modifiedShape, modifiedStrides); |
| rewriter.replaceOp(createNdOp, newCreateNdDescOp.getResult()); |
| return success(); |
| } |
| }; |
| |
| /// Checks if a LoadNdOp consumes a tensor desc type that was rewritten for |
| /// tranpose optimization. If so, rewrites the LoadNdOp to to align with the |
| /// adjusted tensor desc type. This can result in multiple LoadNdOps being |
| /// generated to fill in the original load shape. |
| class XeGPULoadNdDescOpPattern final |
| : public OpConversionPattern<xegpu::LoadNdOp> { |
| public: |
| using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern; |
| LogicalResult |
| matchAndRewrite(xegpu::LoadNdOp loadNdOp, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto origTensorDescType = loadNdOp.getTensorDescType(); |
| auto adaptorType = |
| cast<xegpu::TensorDescType>(adaptor.getTensorDesc().getType()); |
| if (adaptorType == origTensorDescType) |
| return failure(); |
| // Offsets must be adjusted based on innerLaneData. |
| auto laneData = getMaybeLaneData(loadNdOp.getTensorDescType()).value(); |
| int64_t innerLaneData = laneData[1]; |
| auto offsets = loadNdOp.getMixedOffsets(); |
| if (offsets.empty()) |
| return rewriter.notifyMatchFailure(loadNdOp, |
| "Expecting offsets in LoadNd"); |
| SmallVector<OpFoldResult> modifiedOffsets(offsets); |
| modifiedOffsets.back() = divideByConstant( |
| rewriter, loadNdOp.getLoc(), |
| convertToValue(rewriter, loadNdOp.getLoc(), modifiedOffsets.back()), |
| innerLaneData); |
| // Get the 2D data shape of this loadNdOp in its original type including |
| // array length. |
| SmallVector<int64_t> origDataShape(origTensorDescType.getShape()); |
| // Adjust the data shape based on innerLaneData. |
| origDataShape.back() /= innerLaneData; |
| // HW supported shape is the new tensor desc shape after conversion. |
| SmallVector<int64_t> hwSupportedShape(adaptorType.getShape()); |
| VectorType origVectorType = |
| VectorType::get(origDataShape, adaptorType.getElementType()); |
| Value data; |
| // Orig data shape is 3D for the array length case. |
| if (origTensorDescType.getArrayLength() > 1) { |
| SmallVector<Value> arraySlices; |
| for (int64_t i = 0; i < origTensorDescType.getArrayLength(); ++i) { |
| Value slice = arith::ConstantOp::create( |
| rewriter, loadNdOp->getLoc(), origVectorType, |
| rewriter.getZeroAttr(origVectorType)); |
| // Increase the Y offset for each array slice. |
| Value offsetY = convertToValue(rewriter, loadNdOp->getLoc(), |
| modifiedOffsets.back()); |
| modifiedOffsets.back() = |
| arith::AddIOp::create( |
| rewriter, loadNdOp->getLoc(), offsetY, |
| arith::ConstantIndexOp::create(rewriter, loadNdOp->getLoc(), |
| i * origDataShape[1]) |
| .getResult()) |
| .getResult(); |
| slice = generateLoads( |
| rewriter, cast<TypedValue<VectorType>>(slice), modifiedOffsets, |
| cast<TypedValue<xegpu::TensorDescType>>(adaptor.getTensorDesc()), |
| loadNdOp); |
| // BitCast back to original load shape without array length. |
| auto bitcastType = VectorType::get(origTensorDescType.getShape(), |
| origTensorDescType.getElementType()); |
| auto bitCastOp = vector::BitCastOp::create(rewriter, loadNdOp->getLoc(), |
| bitcastType, slice); |
| // BitCastOp must have the same layout as the original loadNdOp. |
| xegpu::setTemporaryLayout(bitCastOp->getOpResult(0), |
| origTensorDescType.getLayoutAttr()); |
| arraySlices.push_back(bitCastOp.getResult()); |
| } |
| rewriter.replaceOpWithMultiple(loadNdOp, {arraySlices}); |
| return success(); |
| } |
| data = arith::ConstantOp::create( |
| rewriter, loadNdOp->getLoc(), |
| VectorType::get(origDataShape, adaptorType.getElementType()), |
| rewriter.getZeroAttr(origVectorType)); |
| data = generateLoads( |
| rewriter, cast<TypedValue<VectorType>>(data), modifiedOffsets, |
| cast<TypedValue<xegpu::TensorDescType>>(adaptor.getTensorDesc()), |
| loadNdOp); |
| auto bitCastOp = vector::BitCastOp::create(rewriter, loadNdOp->getLoc(), |
| loadNdOp.getType(), data); |
| // BitCastOp must have the same layout as the original loadNdOp. |
| xegpu::setTemporaryLayout(bitCastOp->getOpResult(0), |
| origTensorDescType.getLayoutAttr()); |
| rewriter.replaceOp(loadNdOp, bitCastOp); |
| return success(); |
| } |
| }; |
| |
| /// Vector ExtractOp must be processed if the original tensor desc type has |
| /// array length greater than 1. In this case, the LoadNdOp is replaced with |
| /// multiple LoadNdOps for each array slice making the extraction unnecessary. |
| /// In this case, we simply remove the ExtractOp. |
| class VectorExtractOpPattern final |
| : public OpConversionPattern<vector::ExtractOp> { |
| public: |
| using OpConversionPattern<vector::ExtractOp>::OpConversionPattern; |
| LogicalResult |
| matchAndRewrite(vector::ExtractOp extractOp, OneToNOpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| // Check if the source of the extraction is split to multiple values. |
| if (adaptor.getSource().size() == 1) |
| return failure(); |
| auto mixedPos = extractOp.getMixedPosition(); |
| if (mixedPos.size() != 1) |
| return failure(); |
| auto mayBeInt = getConstantIntValue(mixedPos[0]); |
| if (!mayBeInt) |
| return failure(); |
| rewriter.replaceOp(extractOp, adaptor.getSource()[*mayBeInt]); |
| return success(); |
| } |
| }; |
| |
| /// Performs a reduction over 2 dimensions by decomposing it into two 1D |
| /// reductions ordered based on layout to minimize cross-lane communication. |
| class MultiRed2dOpPattern |
| : public OpConversionPattern<vector::MultiDimReductionOp> { |
| using OpConversionPattern::OpConversionPattern; |
| LogicalResult |
| matchAndRewrite(vector::MultiDimReductionOp reductionOp, OpAdaptor adaptor, |
| ConversionPatternRewriter &rewriter) const override { |
| auto sourceVecType = reductionOp.getSourceVectorType(); |
| if (reductionOp.getReductionDims().size() != 2) |
| return rewriter.notifyMatchFailure(reductionOp, "Expected 2D reduction"); |
| auto resLayout = xegpu::getDistributeLayoutAttr(reductionOp.getResult()); |
| // Retrieve and order dims for 1D decomposition (prefer intra-lane first). |
| auto dims = llvm::to_vector(reductionOp.getReductionDims()); |
| auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, resLayout); |
| // Order does not matter |
| if (intraLaneDim == -1 || crossLaneDim == -1) { |
| intraLaneDim = dims[0]; |
| crossLaneDim = dims[1]; |
| } |
| auto loc = reductionOp.getLoc(); |
| auto acc = reductionOp.getAcc(); |
| |
| // The decomposition below splits the 2D reduction into an intra-lane |
| // then a cross-lane 1D reduction. The natural result layout of the |
| // decomposed sequence (a doubly-sliced layout) differs from the |
| // original 2D reduction's result layout that the rest of the IR was |
| // written/propagated against. To keep the post-peephole IR |
| // self-consistent without depending on a follow-up layout |
| // propagation pass, we always insert a bridge xegpu.convert_layout |
| // from the natural post-decomposition layout to the original |
| // reduction's result layout. Trivial bridges fold away in |
| // canonicalization. |
| xegpu::DistributeLayoutAttr postDecompLayout; |
| if (resLayout) { |
| // Derive the source vector's layout. |
| xegpu::DistributeLayoutAttr srcLayoutForCvt; |
| if (auto resSlice = dyn_cast_if_present<xegpu::SliceAttr>(resLayout)) |
| srcLayoutForCvt = resSlice.getParent(); |
| if (!srcLayoutForCvt) |
| srcLayoutForCvt = |
| xegpu::getDistributeLayoutAttr(reductionOp.getSource()); |
| if (srcLayoutForCvt) { |
| // The natural layout of the post-decomposition reduction result |
| // is a nested SliceAttr: REDUCE_1 (reduces `intraLaneDim` from |
| // the source) yields `slice<src, [intraLaneDim]>`; REDUCE_2 |
| // then reduces `adjCrossLaneDim` from that intermediate, giving |
| // `slice<slice<src, [intraLaneDim]>, [adjCrossLaneDim]>`. |
| MLIRContext *ctx = reductionOp.getContext(); |
| int64_t adjCrossLaneDim = |
| crossLaneDim > intraLaneDim ? crossLaneDim - 1 : crossLaneDim; |
| auto intermediateLayout = xegpu::SliceAttr::get( |
| ctx, srcLayoutForCvt, DenseI64ArrayAttr::get(ctx, {intraLaneDim})); |
| postDecompLayout = xegpu::SliceAttr::get( |
| ctx, intermediateLayout, |
| DenseI64ArrayAttr::get(ctx, {adjCrossLaneDim})); |
| } |
| } |
| |
| SmallVector<int64_t> accShape(sourceVecType.getShape()); |
| accShape.erase(accShape.begin() + intraLaneDim); |
| Type eTy = sourceVecType.getElementType(); |
| Value constNeutralVal = xegpu::createReductionNeutralValue( |
| rewriter, loc, VectorType::get(accShape, eTy), reductionOp.getKind()); |
| |
| Value intraLaneReduced = vector::MultiDimReductionOp::create( |
| rewriter, loc, reductionOp.getKind(), reductionOp.getSource(), |
| constNeutralVal, ArrayRef<int64_t>(intraLaneDim)); |
| |
| // Adjust crossLaneDim after the first reduction. |
| if (crossLaneDim > intraLaneDim) |
| crossLaneDim -= 1; |
| Value crossLaneReduced = vector::MultiDimReductionOp::create( |
| rewriter, loc, reductionOp.getKind(), intraLaneReduced, acc, |
| ArrayRef<int64_t>(crossLaneDim)); |
| assert(crossLaneReduced.getType() == reductionOp.getResult().getType() && |
| "Type mismatch"); |
| |
| Value replacement = crossLaneReduced; |
| if (resLayout && postDecompLayout) { |
| // Bridge from the natural post-decomposition layout to the |
| // original reduction's result layout. This preserves the contract |
| // any consumer (convert_layout, anchor op, or otherwise) was |
| // written against, so the rewrite is correct independent of |
| // whether layout propagation runs afterwards. |
| auto bridgeOp = xegpu::ConvertLayoutOp::create( |
| rewriter, loc, crossLaneReduced.getType(), crossLaneReduced, |
| postDecompLayout, resLayout); |
| replacement = bridgeOp.getResult(); |
| } |
| |
| rewriter.replaceOp(reductionOp, replacement); |
| return success(); |
| } |
| |
| private: |
| std::pair<int64_t, int64_t> |
| getReductionDimOrder(ArrayRef<int64_t> reductionDims, |
| xegpu::DistributeLayoutAttr layout) const { |
| assert(layout.isForSubgroup() && "Must know the lane layout"); |
| assert(reductionDims.size() == 2 && "Expected 2D reduction"); |
| int64_t intra, cross = -1; |
| xegpu::LayoutAttr layoutAttr = dyn_cast<xegpu::LayoutAttr>(layout); |
| if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout)) |
| layoutAttr = |
| dyn_cast<xegpu::LayoutAttr>(layoutSliceAttr.flatten().getParent()); |
| assert(layoutAttr); |
| SmallVector<int64_t> laneLayout = layoutAttr.getEffectiveLaneLayoutAsInt(); |
| |
| assert(laneLayout.size() && "Expected a non-empty layout"); |
| // try to pick a dim that does not communicate |
| for (auto dim : reductionDims) { |
| if (laneLayout[dim] == 1) |
| intra = dim; |
| else |
| cross = dim; |
| } |
| return {intra, cross}; |
| } |
| }; |
| |
| } // namespace |
| |
| void xegpu::populateXeGPUPeepHoleOptimizerPatterns( |
| RewritePatternSet &patterns) { |
| patterns.add<XeGPUCreateNdDescOpPattern, XeGPULoadNdDescOpPattern, |
| VectorExtractOpPattern, MultiRed2dOpPattern>( |
| patterns.getContext()); |
| } |
| |
| namespace { |
| |
| struct XeGPUPeepHoleOptimizerPass final |
| : public xegpu::impl::XeGPUPeepHoleOptimizerBase< |
| XeGPUPeepHoleOptimizerPass> { |
| void runOnOperation() override { |
| MLIRContext &context = getContext(); |
| TypeConverter converter; |
| RewritePatternSet patterns(&context); |
| ConversionTarget target(context); |
| |
| // This pass is only meant for PVC, BMG or CRI targets. If unsupported |
| // target is found, exit early. |
| bool isTargetSupported = false; |
| getOperation()->walk([&](gpu::GPUFuncOp funcOp) { |
| auto chipStr = xegpu::getChipStr(funcOp); |
| if (chipStr && (chipStr.value() == "pvc" || chipStr.value() == "bmg" || |
| chipStr.value() == "cri")) |
| isTargetSupported = true; |
| }); |
| |
| if (!isTargetSupported) { |
| DBGS() << "XeGPUPeepHoleOptimizerPass only supports PVC, BMG targets." |
| << "\n"; |
| return; |
| } |
| |
| // Run array length optimization patterns first so that subsequent transpose |
| // peephole patterns operate on the array-length-optimized tensor descs. |
| { |
| RewritePatternSet arrayLenPatterns(&context); |
| xegpu::populateXeGPUArrayLengthOptimizationPatterns(arrayLenPatterns); |
| if (failed(applyPatternsGreedily(getOperation(), |
| std::move(arrayLenPatterns)))) { |
| DBGS() << "Array length optimization patterns failed.\n"; |
| return signalPassFailure(); |
| } |
| } |
| |
| // CreateNdDescOp and LoadNdOp with optimizable tensor desc types must be |
| // converted. |
| target.addDynamicallyLegalOp<xegpu::CreateNdDescOp>( |
| [&](xegpu::CreateNdDescOp createNdOp) { |
| return !canBeOptimizedForTranspose(createNdOp.getType()); |
| }); |
| target.addDynamicallyLegalOp<xegpu::LoadNdOp>( |
| [&](xegpu::LoadNdOp loadNdOp) { |
| return !canBeOptimizedForTranspose(loadNdOp.getTensorDescType()); |
| }); |
| // Vector ExtractOps can have optimizable layouts if they extract from |
| // LoadNdOps with array length greater than 1. These ExtractOps must be |
| // converted. |
| target.addDynamicallyLegalOp<vector::ExtractOp>( |
| [&](vector::ExtractOp extractOp) { |
| auto layout = xegpu::getTemporaryLayout( |
| dyn_cast<OpResult>(extractOp.getResult())); |
| if (!layout) |
| return true; |
| auto laneLayout = layout.getEffectiveLaneLayoutAsInt(); |
| auto laneData = layout.getEffectiveLaneDataAsInt(); |
| return !canBeOptimizedForTranspose(laneLayout, laneData); |
| }); |
| |
| target.addDynamicallyLegalOp<vector::MultiDimReductionOp>( |
| [=](Operation *op) -> bool { |
| auto layout = xegpu::getDistributeLayoutAttr(op->getResult(0)); |
| if (!layout || !layout.isForSubgroup()) |
| return true; |
| if (auto reductionOp = dyn_cast<vector::MultiDimReductionOp>(op)) |
| return reductionOp.getReductionDims().size() != 2; |
| return true; |
| }); |
| |
| converter.addConversion([](Type type) { return type; }); |
| |
| target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect, |
| vector::VectorDialect>(); |
| // xegpu.convert_layout is left untouched by this pass; mark it legal |
| // so in-place updates don't trigger re-legalization failures. |
| target.addLegalOp<xegpu::ConvertLayoutOp>(); |
| scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns, |
| target); |
| xegpu::populateXeGPUPeepHoleOptimizerPatterns(patterns); |
| if (failed(applyPartialConversion(getOperation(), target, |
| std::move(patterns)))) { |
| DBGS() << "Optimize block loads pass failed.\n"; |
| return signalPassFailure(); |
| } |
| |
| // Apply folding for cleaning up IR. |
| MLIRContext *ctx = &getContext(); |
| RewritePatternSet emptyPatterns(ctx); |
| (void)applyPatternsGreedily(getOperation(), std::move(emptyPatterns)); |
| |
| xegpu::removeTemporaryLayoutAttrs(getOperation()); |
| } |
| }; |
| |
| } // namespace |
