112 lines · cpp
1//===-- XeGPUVectorLinearize.cpp - Linearizes n-D vectors to 1-D vectors --===//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/SCF/IR/SCF.h"10#include "mlir/Dialect/SCF/Transforms/Patterns.h"11#include "mlir/Dialect/Vector/IR/VectorOps.h"12#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"13#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"14#include "mlir/Dialect/XeGPU/IR/XeGPU.h"15#include "mlir/Dialect/XeGPU/Transforms/Passes.h"16#include "mlir/Pass/Pass.h"17#include "mlir/Transforms/DialectConversion.h"18#include "mlir/Transforms/GreedyPatternRewriteDriver.h"19#include "llvm/Support/Debug.h"20#include "llvm/Support/DebugLog.h"21#include "llvm/Support/raw_ostream.h"22 23#include <optional>24 25namespace mlir {26namespace xegpu {27#define GEN_PASS_DEF_XEGPUVECTORLINEARIZE28#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"29} // namespace xegpu30} // namespace mlir31 32#define DEBUG_TYPE "xegpu-vector-linearize"33 34using namespace mlir;35 36namespace {37struct XeGPUVectorLinearizePass final38 : public xegpu::impl::XeGPUVectorLinearizeBase<XeGPUVectorLinearizePass> {39 void runOnOperation() override {40 // vector.broadcast and vector.gather requires progressive lowering41 {42 RewritePatternSet patterns(&getContext());43 vector::populateVectorBroadcastLoweringPatterns(patterns);44 vector::populateVectorGatherLoweringPatterns(patterns);45 vector::populateVectorGatherToConditionalLoadPatterns(patterns);46 // vector.transpose lowering47 // Shuffle16x16 will fallback to Shuffle1D for non 16x16 sizes.48 vector::populateVectorTransposeLoweringPatterns(49 patterns, vector::VectorTransposeLowering::Shuffle16x16);50 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))51 return signalPassFailure();52 }53 54 // Unroll load/store from <d1xd2x...xdk> to (d1*d2*...*d(k-1)) slices of55 // <1x1x...x1xdk>.56 {57 RewritePatternSet patterns(&getContext());58 vector::UnrollVectorOptions vectorOptions;59 vectorOptions.setNativeShapeFn(60 [](Operation *op) -> std::optional<SmallVector<int64_t>> {61 auto extractVectorType = [](Operation *op) -> VectorType {62 if (auto loadOp = dyn_cast<vector::LoadOp>(op))63 return loadOp.getVectorType();64 if (auto storeOp = dyn_cast<vector::StoreOp>(op))65 return storeOp.getVectorType();66 return nullptr;67 };68 69 VectorType vecType = extractVectorType(op);70 if (!vecType)71 return std::nullopt;72 73 // Only handle rank >= 2 so we actually unroll something.74 int64_t rank = vecType.getRank();75 if (rank < 2)76 return std::nullopt;77 78 ArrayRef<int64_t> shape = vecType.getShape();79 // Produce native shape: 1 x 1 x ... x (original last dim).80 SmallVector<int64_t> native(rank, 1);81 native.back() = shape.back();82 return native;83 });84 vector::populateVectorUnrollPatterns(patterns, vectorOptions);85 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {86 LDBG() << "Unroll failed.";87 return signalPassFailure();88 }89 }90 91 // Use vector linearization patterns92 {93 MLIRContext &context = getContext();94 TypeConverter converter;95 RewritePatternSet patterns(&context);96 ConversionTarget target(context);97 vector::populateForVectorLinearize(converter, target);98 vector::populateVectorLinearizeBasePatterns(converter, target, patterns);99 vector::populateVectorLinearizeShuffleLikeOpsPatterns(converter, target,100 patterns);101 scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,102 target);103 if (failed(applyPartialConversion(getOperation(), target,104 std::move(patterns)))) {105 LDBG() << "Linearization failed.";106 return signalPassFailure();107 }108 }109 }110};111} // namespace112