118 lines · cpp
1//===- TestAllReduceLowering.cpp - Test gpu.all_reduce lowering -----------===//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// This file contains test passes for lowering the gpu.all_reduce op.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"14#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"15#include "mlir/Dialect/Arith/IR/Arith.h"16#include "mlir/Dialect/Func/IR/FuncOps.h"17#include "mlir/Dialect/GPU/Transforms/Passes.h"18#include "mlir/Dialect/Index/IR/IndexDialect.h"19#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"20#include "mlir/Dialect/MemRef/IR/MemRef.h"21#include "mlir/Dialect/Vector/IR/VectorOps.h"22#include "mlir/IR/PatternMatch.h"23#include "mlir/Pass/Pass.h"24#include "mlir/Transforms/GreedyPatternRewriteDriver.h"25 26using namespace mlir;27 28namespace {29struct TestGpuRewritePass30 : public PassWrapper<TestGpuRewritePass, OperationPass<ModuleOp>> {31 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestGpuRewritePass)32 33 void getDependentDialects(DialectRegistry ®istry) const override {34 registry.insert<arith::ArithDialect, func::FuncDialect, index::IndexDialect,35 memref::MemRefDialect>();36 }37 StringRef getArgument() const final { return "test-gpu-rewrite"; }38 StringRef getDescription() const final {39 return "Applies all rewrite patterns within the GPU dialect.";40 }41 void runOnOperation() override {42 RewritePatternSet patterns(&getContext());43 populateGpuRewritePatterns(patterns);44 populateGpuSubgroupIdPatterns(patterns);45 (void)applyPatternsGreedily(getOperation(), std::move(patterns));46 }47};48 49struct TestGpuSubgroupReduceLoweringPass50 : public PassWrapper<TestGpuSubgroupReduceLoweringPass,51 OperationPass<ModuleOp>> {52 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(53 TestGpuSubgroupReduceLoweringPass)54 55 TestGpuSubgroupReduceLoweringPass() = default;56 TestGpuSubgroupReduceLoweringPass(57 const TestGpuSubgroupReduceLoweringPass &pass)58 : PassWrapper(pass) {}59 60 void getDependentDialects(DialectRegistry ®istry) const override {61 registry62 .insert<amdgpu::AMDGPUDialect, arith::ArithDialect, LLVM::LLVMDialect,63 ROCDL::ROCDLDialect, vector::VectorDialect>();64 }65 66 StringRef getArgument() const final {67 return "test-gpu-subgroup-reduce-lowering";68 }69 70 StringRef getDescription() const final {71 return "Applies gpu.subgroup_reduce lowering patterns.";72 }73 74 Option<bool> expandToShuffles{75 *this, "expand-to-shuffles",76 llvm::cl::desc("Expand subgroup_reduce ops to shuffle ops."),77 llvm::cl::init(false)};78 79 Option<std::string> target{80 *this, "target",81 llvm::cl::desc("Target backend name which will be used to provide "82 "compatible lowerings of subgroup reduce."),83 llvm::cl::init("")};84 85 void runOnOperation() override {86 RewritePatternSet patterns(&getContext());87 88 // Since both pattern sets match on the same ops, set higher benefit to89 // perform fewer failing matches.90 populateGpuBreakDownSubgroupReducePatterns(patterns,91 /*maxShuffleBitwidth=*/32,92 PatternBenefit(3));93 if (expandToShuffles) {94 auto maybeChipset = amdgpu::Chipset::parse(target);95 if (succeeded(maybeChipset)) {96 populateGpuLowerSubgroupReduceToDPPPatterns(97 patterns, /*subgroupSize=*/64, *maybeChipset, PatternBenefit(2));98 populateGpuLowerClusteredSubgroupReduceToDPPPatterns(99 patterns, /*subgroupSize=*/64, *maybeChipset, PatternBenefit(2));100 }101 populateGpuLowerSubgroupReduceToShufflePatterns(102 patterns, /*subgroupSize=*/32, /*shuffleBitwidth=*/32);103 populateGpuLowerClusteredSubgroupReduceToShufflePatterns(104 patterns, /*subgroupSize=*/32, /*shuffleBitwidth=*/32);105 }106 107 (void)applyPatternsGreedily(getOperation(), std::move(patterns));108 }109};110} // namespace111 112namespace mlir {113void registerTestGpuLoweringPasses() {114 PassRegistration<TestGpuRewritePass>();115 PassRegistration<TestGpuSubgroupReduceLoweringPass>();116}117} // namespace mlir118