350 lines · cpp
1//===- TestXeGPUTransforms.cpp -- Test Vector transforms and lowerings ----===//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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/GPU/IR/GPUDialect.h"10#include "mlir/Dialect/Index/IR/IndexDialect.h"11#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"12#include "mlir/Dialect/XeGPU/IR/XeGPU.h"13#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"14#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"15#include "mlir/Pass/Pass.h"16#include "mlir/Pass/PassManager.h"17#include "mlir/Transforms/DialectConversion.h"18#include "mlir/Transforms/GreedyPatternRewriteDriver.h"19 20using namespace mlir;21using namespace mlir::xegpu;22 23namespace {24 25#define DEBUG_TYPE "test-xegpu-unroll"26 27struct TestXeGPUUnrollingPatterns28 : public PassWrapper<TestXeGPUUnrollingPatterns,29 OperationPass<gpu::GPUModuleOp>> {30 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUUnrollingPatterns)31 32 StringRef getArgument() const final {33 return "test-xegpu-unrolling-patterns";34 }35 36 StringRef getDescription() const final {37 return "Test lowering patterns to unroll ops in the xegpu dialect";38 }39 40 void getDependentDialects(::mlir::DialectRegistry ®istry) const override {41 registry.insert<memref::MemRefDialect>();42 registry.insert<xegpu::XeGPUDialect>();43 registry.insert<vector::VectorDialect>();44 }45 46 TestXeGPUUnrollingPatterns() = default;47 TestXeGPUUnrollingPatterns(const TestXeGPUUnrollingPatterns &pass)48 : PassWrapper(pass) {}49 50 void runOnOperation() override {51 MLIRContext *ctx = &getContext();52 xegpu::UnrollOptions options;53 options.setNativeShapeFn([&](Operation *op)54 -> std::optional<SmallVector<int64_t>> {55 if (isa<xegpu::CreateNdDescOp, xegpu::UpdateNdOffsetOp,56 xegpu::PrefetchNdOp, xegpu::LoadNdOp, xegpu::StoreNdOp,57 xegpu::CreateDescOp, xegpu::UpdateOffsetOp, xegpu::PrefetchOp,58 xegpu::LoadGatherOp, xegpu::StoreScatterOp>(op)) {59 xegpu::TensorDescType tdescTy;60 if (auto createNdOp = dyn_cast<xegpu::CreateNdDescOp>(op)) {61 tdescTy = createNdOp.getType();62 } else if (auto updateNdOp = dyn_cast<xegpu::UpdateNdOffsetOp>(op)) {63 tdescTy = updateNdOp.getTensorDescType();64 } else if (auto prefetchNdOp = dyn_cast<xegpu::PrefetchNdOp>(op)) {65 tdescTy = prefetchNdOp.getTensorDescType();66 } else if (auto loadNdOp = dyn_cast<xegpu::LoadNdOp>(op)) {67 tdescTy = loadNdOp.getTensorDescType();68 } else if (auto storeNdOp = dyn_cast<xegpu::StoreNdOp>(op)) {69 tdescTy = storeNdOp.getTensorDescType();70 } else if (auto createOp = dyn_cast<xegpu::CreateDescOp>(op)) {71 tdescTy = createOp.getType();72 } else if (auto updateOp = dyn_cast<xegpu::UpdateOffsetOp>(op)) {73 tdescTy = updateOp.getTensorDescType();74 } else if (auto prefetchOp = dyn_cast<xegpu::PrefetchOp>(op)) {75 tdescTy = prefetchOp.getTensorDescType();76 } else if (auto loadOp = dyn_cast<xegpu::LoadGatherOp>(op)) {77 if (loadOp.getOffsets()) {78 auto layout = xegpu::getDistributeLayoutAttr(loadOp.getResult());79 if (layout && layout.isForSubgroup()) {80 auto inst_data = layout.getEffectiveInstDataAsInt();81 if (!inst_data.empty())82 return SmallVector<int64_t>(inst_data.begin(), inst_data.end());83 }84 return std::nullopt;85 }86 tdescTy = loadOp.getTensorDescType();87 } else if (auto storeOp = dyn_cast<xegpu::StoreScatterOp>(op)) {88 if (storeOp.getOffsets()) {89 auto layout = llvm::dyn_cast_or_null<xegpu::LayoutAttr>(90 op->getAttr("layout"));91 if (layout && layout.isForSubgroup()) {92 auto inst_data = layout.getEffectiveInstDataAsInt();93 if (!inst_data.empty())94 return SmallVector<int64_t>(inst_data.begin(), inst_data.end());95 }96 return std::nullopt;97 }98 tdescTy = storeOp.getTensorDescType();99 }100 101 if (auto layout = tdescTy.getLayoutAttr()) {102 auto inst_data = layout.getInstData();103 if (inst_data && layout.isForSubgroup())104 return SmallVector<int64_t>(inst_data.asArrayRef().begin(),105 inst_data.asArrayRef().end());106 }107 }108 109 if (isa<xegpu::DpasOp>(op))110 return SmallVector<int64_t>{8, 16, 16};111 112 return std::nullopt;113 });114 115 options.setUnrolledTypesFn(116 [&](ShapedType type, ArrayRef<int64_t> tileShape,117 bool returnSingleType = false) -> SmallVector<Type> {118 Type elemTy = type.getElementType();119 Type newTy;120 121 // TensorDescType needs to drop the inst_data field in the layout122 // attribute123 if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(type)) {124 Attribute encoding = tdescTy.getEncoding();125 auto layout = tdescTy.getLayoutAttr();126 127 // If the encoding is a ScatterTensorDescAttr, we need to128 // potentially adjust the chunk size based on the inst_data.129 if (tdescTy.isScattered()) {130 int64_t chunkSize = tdescTy.getChunkSizeAsInt();131 132 if (chunkSize > 1) {133 int64_t blockedChunkSize = chunkSize;134 auto instData = layout.getInstData();135 if (!instData.empty())136 blockedChunkSize = instData.asArrayRef().back();137 138 // To create a new attribute with a different chunk_size:139 auto newEncoding = xegpu::ScatterTensorDescAttr::get(140 ctx, tdescTy.getMemorySpace(), blockedChunkSize);141 142 encoding = newEncoding;143 }144 }145 if (layout) {146 if (layout.getLaneLayout() == nullptr)147 layout = xegpu::LayoutAttr();148 else149 layout = layout.dropInstData();150 }151 152 newTy = xegpu::TensorDescType::get(ctx, tileShape, elemTy, encoding,153 layout);154 155 } else {156 newTy = type.clone(tileShape, elemTy);157 }158 159 if (returnSingleType)160 return SmallVector<Type>{newTy};161 std::optional<SmallVector<int64_t>> ratio =162 computeShapeRatio(type.getShape(), tileShape);163 assert(ratio && "Expecting the ratio to be valid.");164 return SmallVector<Type>(computeProduct(*ratio), newTy);165 });166 167 RewritePatternSet patterns(ctx);168 169 populateXeGPUUnrollPatterns(patterns, options);170 (void)applyPatternsGreedily(getOperation(), std::move(patterns));171 }172};173 174#undef DEBUG_TYPE175#define DEBUG_TYPE "test-xegpu-layout-interface"176 177// Test pattern for distributing vector::StepOp from workgroup to subgroup.178// Validates DistributeLayoutAttr interfaces for offset computation179// abstraction between LayoutAttr and SliceAttr.180class TestStepOpPattern : public OpConversionPattern<vector::StepOp> {181 using OpConversionPattern<vector::StepOp>::OpConversionPattern;182 183 LogicalResult184 matchAndRewrite(vector::StepOp op, OneToNOpAdaptor adaptor,185 ConversionPatternRewriter &rewriter) const override {186 187 auto layoutName = xegpu::getLayoutName(op->getResult(0));188 auto sliceAttr = op->getAttrOfType<xegpu::SliceAttr>(layoutName);189 if (!sliceAttr || sliceAttr.getRank() != 1)190 return failure();191 192 std::optional<SmallVector<int64_t>> sgShape =193 sliceAttr.getEffectiveSgDataAsInt();194 if (!sgShape)195 return failure();196 197 Location loc = op.getLoc();198 VectorType type = op.getResult().getType();199 auto wgShape = type.getShape();200 201 Value sgId =202 gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);203 auto maybeOffsets =204 sliceAttr.computeDistributedCoords(rewriter, loc, sgId, wgShape);205 if (failed(maybeOffsets))206 return failure();207 208 VectorType newTy = type.cloneWith(*sgShape, type.getElementType());209 Value base = vector::StepOp::create(rewriter, loc, newTy);210 SmallVector<Value> newOps;211 for (auto offsets : *maybeOffsets) {212 Value bcast =213 vector::BroadcastOp::create(rewriter, loc, newTy, offsets[0]);214 Value add = arith::AddIOp::create(rewriter, loc, base, bcast);215 newOps.push_back(add);216 }217 rewriter.replaceOpWithMultiple(op, {newOps});218 return success();219 }220};221 222struct TestXeGPUSGDistribute223 : public PassWrapper<TestXeGPUSGDistribute,224 OperationPass<gpu::GPUModuleOp>> {225 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUSGDistribute)226 227 StringRef getArgument() const final { return "test-xegpu-sg-distribute"; }228 229 StringRef getDescription() const final {230 return "Test the implementation of XeGPU Subgroup Distribution";231 }232 233 void getDependentDialects(::mlir::DialectRegistry ®istry) const override {234 registry.insert<arith::ArithDialect>();235 registry.insert<memref::MemRefDialect>();236 registry.insert<xegpu::XeGPUDialect>();237 registry.insert<vector::VectorDialect>();238 registry.insert<index::IndexDialect>();239 }240 241 TestXeGPUSGDistribute() = default;242 TestXeGPUSGDistribute(const TestXeGPUSGDistribute &pass) = default;243 244 void runOnOperation() override {245 RewritePatternSet patterns(&getContext());246 xegpu::populateXeGPUSubgroupDistributePatterns(patterns);247 (void)applyPatternsGreedily(getOperation(), std::move(patterns));248 }249};250 251struct TestXeGPUMoveFuncBodyToWarpOp252 : public PassWrapper<TestXeGPUMoveFuncBodyToWarpOp,253 OperationPass<gpu::GPUModuleOp>> {254 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPUMoveFuncBodyToWarpOp)255 256 StringRef getArgument() const final {257 return "test-xegpu-move-func-to-warp-op";258 }259 260 StringRef getDescription() const final {261 return "Test the implementation of XeGPU move gpu function body to "262 "WarpExecuteOnLane0 op.";263 }264 265 void getDependentDialects(::mlir::DialectRegistry ®istry) const override {266 registry.insert<xegpu::XeGPUDialect>();267 registry.insert<gpu::GPUDialect>();268 }269 270 TestXeGPUMoveFuncBodyToWarpOp() = default;271 TestXeGPUMoveFuncBodyToWarpOp(const TestXeGPUMoveFuncBodyToWarpOp &pass) =272 default;273 274 void runOnOperation() override {275 RewritePatternSet patterns(&getContext());276 xegpu::populateXeGPUMoveFuncBodyToWarpOpPatterns(patterns);277 (void)applyPatternsGreedily(getOperation(), std::move(patterns));278 }279};280 281struct TestXeGPULayoutInterface282 : public PassWrapper<TestXeGPULayoutInterface,283 OperationPass<gpu::GPUModuleOp>> {284 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPULayoutInterface)285 286 StringRef getArgument() const final { return "test-xegpu-layout-interface"; }287 288 StringRef getDescription() const final {289 return "Test the implementation of XeGPU Layout interfaces";290 }291 292 void getDependentDialects(::mlir::DialectRegistry ®istry) const override {293 registry.insert<arith::ArithDialect>();294 registry.insert<memref::MemRefDialect>();295 registry.insert<xegpu::XeGPUDialect>();296 registry.insert<vector::VectorDialect>();297 registry.insert<index::IndexDialect>();298 }299 300 TestXeGPULayoutInterface() = default;301 TestXeGPULayoutInterface(const TestXeGPULayoutInterface &pass)302 : PassWrapper(pass) {}303 304 void runOnOperation() override {305 MLIRContext *ctx = &getContext();306 307 TypeConverter typeConverter;308 auto materializeCast = [&](mlir::OpBuilder &builder, mlir::Type type,309 mlir::ValueRange inputs,310 mlir::Location loc) -> mlir::Value {311 return UnrealizedConversionCastOp::create(builder, loc, type, inputs)312 .getResult(0);313 };314 typeConverter.addSourceMaterialization(materializeCast);315 typeConverter.addTargetMaterialization(materializeCast);316 317 RewritePatternSet patterns(ctx);318 patterns.add<TestStepOpPattern>(typeConverter, ctx);319 320 ConversionTarget target(*ctx);321 auto isLegal = [&](xegpu::SliceAttr layout) -> bool {322 return !layout || !layout.isForWorkgroup();323 };324 325 target.addDynamicallyLegalOp<vector::StepOp>(326 [&](vector::StepOp op) -> bool {327 auto layoutName = xegpu::getLayoutName(op->getResult(0));328 auto sliceAttr = op->getAttrOfType<xegpu::SliceAttr>(layoutName);329 return isLegal(sliceAttr);330 });331 332 target.markUnknownOpDynamicallyLegal([](Operation *op) { return true; });333 334 (void)applyPartialConversion(getOperation(), target, std::move(patterns));335 }336};337 338} // namespace339 340namespace mlir {341namespace test {342void registerTestXeGPULowerings() {343 PassRegistration<TestXeGPUUnrollingPatterns>();344 PassRegistration<TestXeGPULayoutInterface>();345 PassRegistration<TestXeGPUSGDistribute>();346 PassRegistration<TestXeGPUMoveFuncBodyToWarpOp>();347}348} // namespace test349} // namespace mlir350