Skip to content

Commit

Permalink
[LLVMGPU] Switch to new pass generation tablegen definitions. (#18213)
Browse files Browse the repository at this point in the history
This is mostly an NFC change. The revision applies a little cleanups:

- Switch a couple of passes to follow `create.*Pass` naming convention.
- Move the detailed comments from `Passes.h` to `Passes.td`.

---------

Signed-off-by: hanhanW <hanhan0912@gmail.com>
  • Loading branch information
hanhanW committed Aug 14, 2024
1 parent 0c2f51b commit a72e78b
Show file tree
Hide file tree
Showing 31 changed files with 224 additions and 389 deletions.
5 changes: 0 additions & 5 deletions compiler/src/iree/compiler/Codegen/Common/Passes.h
Original file line number Diff line number Diff line change
Expand Up @@ -88,11 +88,6 @@ createTileAndDistributeToWorkgroupsPass(
int32_t maxWorkgroupParallelDims,
linalg::DistributionMethod distributionMethod);

// TODO(hanchung): Move it where it is defined (i.e., Codegen/LLVMGPU).
// Extract address computations (including the ones with GPU instructions) into
// their own separate instructions.
std::unique_ptr<Pass> createExtractAddressComputationGPUPass();

//----------------------------------------------------------------------------//
// CodeGen Common Patterns
//----------------------------------------------------------------------------//
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,17 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <numeric>
#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"

#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/VectorOpUtils.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_AMDGPUPREPAREFORCHAINEDMATMULPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

using VectorValue = TypedValue<VectorType>;

namespace {
Expand Down Expand Up @@ -59,8 +62,8 @@ namespace {
/// C = A @ B --> C.T = B.T @ A.T
/// is only defined on standard "@" function, it may be a different
/// transformation for other indexing maps.
struct AMDGPUPrepareForChainedMatmulPass
: public AMDGPUPrepareForChainedMatmulBase<
struct AMDGPUPrepareForChainedMatmulPass final
: impl::AMDGPUPrepareForChainedMatmulPassBase<
AMDGPUPrepareForChainedMatmulPass> {
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<vector::VectorDialect>();
Expand Down Expand Up @@ -255,10 +258,4 @@ struct AMDGPUPrepareForChainedMatmulPass
};

} // namespace

std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
createAMDGPUPrepareForChainedMatmulPass() {
return std::make_unique<AMDGPUPrepareForChainedMatmulPass>();
}

} // namespace mlir::iree_compiler
2 changes: 0 additions & 2 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ iree_gentbl_cc_library(
iree_compiler_cc_library(
name = "PassHeaders",
hdrs = [
"PassDetail.h",
"Passes.h",
"Passes.h.inc",
],
Expand Down Expand Up @@ -69,7 +68,6 @@ iree_gentbl_cc_library(
iree_compiler_cc_library(
name = "ROCDLPassHeaders",
hdrs = [
"ROCDLPassDetail.h",
"ROCDLPasses.h",
"ROCDLPasses.h.inc",
],
Expand Down
2 changes: 0 additions & 2 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@ iree_cc_library(
NAME
PassHeaders
HDRS
"PassDetail.h"
"Passes.h"
"Passes.h.inc"
DEPS
Expand Down Expand Up @@ -52,7 +51,6 @@ iree_cc_library(
NAME
ROCDLPassHeaders
HDRS
"ROCDLPassDetail.h"
"ROCDLPasses.h"
"ROCDLPasses.h.inc"
DEPS
Expand Down
13 changes: 6 additions & 7 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

#include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"

#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
Expand All @@ -25,6 +24,9 @@

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_TESTLLVMGPUSCALARIZEMATHOPPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

void ConvertToDynamicSharedMemory(ModuleOp moduleOp) {
SymbolTableCollection symbolTableCollection;
// Collect all the adressOfOps to static shared memory globals.
Expand Down Expand Up @@ -183,8 +185,9 @@ struct ConvertSharedMemAllocOp : public OpRewritePattern<memref::AllocOp> {

/// Pass to test in dialect transformation used to legalize the IR before
/// convertToNVVM/ConvertToROCDL.
class TestLLVMGPULegalizeOpPass
: public TestLLVMGPUScalarizeMathOpBase<TestLLVMGPULegalizeOpPass> {
class TestLLVMGPULegalizeOpPass final
: public impl::TestLLVMGPUScalarizeMathOpPassBase<
TestLLVMGPULegalizeOpPass> {
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<vector::VectorDialect>();
}
Expand Down Expand Up @@ -542,10 +545,6 @@ void populateLowerHALInterfaceOp(RewritePatternSet &patterns) {
patterns.getContext());
}

std::unique_ptr<OperationPass<ModuleOp>> createTestLLVMGPULegalizePass() {
return std::make_unique<TestLLVMGPULegalizeOpPass>();
}

static IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
return IntegerAttr::get(IntegerType::get(ctx, 64), space);
}
Expand Down
15 changes: 8 additions & 7 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
#include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"
#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
Expand All @@ -24,6 +23,7 @@
#include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
#include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
Expand All @@ -35,14 +35,20 @@

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_CONVERTTONVVMPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

namespace {

/// A pass that replaces all occurrences of GPU device operations with their
/// corresponding NVVM equivalent.
///
/// This pass only handles device code and is not meant to be run on GPU host
/// code.
struct ConvertToNVVMPass : public ConvertToNVVMBase<ConvertToNVVMPass> {
struct ConvertToNVVMPass final
: impl::ConvertToNVVMPassBase<ConvertToNVVMPass> {
using impl::ConvertToNVVMPassBase<ConvertToNVVMPass>::ConvertToNVVMPassBase;

void getDependentDialects(DialectRegistry &registry) const override {
registry
.insert<gpu::GPUDialect, IREE::GPU::IREEGPUDialect, LLVM::LLVMDialect,
Expand Down Expand Up @@ -180,9 +186,4 @@ struct ConvertToNVVMPass : public ConvertToNVVMBase<ConvertToNVVMPass> {
};

} // namespace

std::unique_ptr<OperationPass<ModuleOp>> createConvertToNVVMPass() {
return std::make_unique<ConvertToNVVMPass>();
}

} // namespace mlir::iree_compiler
15 changes: 8 additions & 7 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToROCDL.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
#include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"
#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h"
Expand Down Expand Up @@ -40,6 +39,9 @@

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_CONVERTTOROCDLPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

static llvm::cl::opt<int>
clROCMIndexingBits("iree-rocm-index-bits",
llvm::cl::desc("Set the bit width of indices in ROCm."),
Expand Down Expand Up @@ -75,7 +77,11 @@ static void populateConvertGPUToAMDGPUPatterns(RewritePatternSet &patterns) {
///
/// This pass only handles device code and is not meant to be run on GPU host
/// code.
struct ConvertToROCDLPass : public ConvertToROCDLBase<ConvertToROCDLPass> {
struct ConvertToROCDLPass final
: impl::ConvertToROCDLPassBase<ConvertToROCDLPass> {
using impl::ConvertToROCDLPassBase<
ConvertToROCDLPass>::ConvertToROCDLPassBase;

void getDependentDialects(DialectRegistry &registry) const override {
registry
.insert<IREE::GPU::IREEGPUDialect, LLVM::LLVMDialect,
Expand Down Expand Up @@ -203,9 +209,4 @@ struct ConvertToROCDLPass : public ConvertToROCDLBase<ConvertToROCDLPass> {
LDBG("After converting to dynamic shared memory\n" << m);
}
};

std::unique_ptr<OperationPass<ModuleOp>> createConvertToROCDLPass() {
return std::make_unique<ConvertToROCDLPass>();
}

} // namespace mlir::iree_compiler
Original file line number Diff line number Diff line change
Expand Up @@ -5,20 +5,22 @@
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Codegen/Common/ExtractAddressComputation.h"
#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "extract-address-computation-gpu"

using namespace mlir;

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_EXTRACTADDRESSCOMPUTATIONGPUPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

//===----------------------------------------------------------------------===//
// Helper functions for the `load base[off0...]`
// => `load (subview base[off0...])[0...]` pattern.
Expand Down Expand Up @@ -80,8 +82,8 @@ populateExtractAddressComputationGPUPatterns(RewritePatternSet &patterns) {
// Pass registration
//===----------------------------------------------------------------------===//
namespace {
struct ExtractAddressComputationGPUPass
: public ExtractAddressComputationGPUBase<
struct ExtractAddressComputationGPUPass final
: impl::ExtractAddressComputationGPUPassBase<
ExtractAddressComputationGPUPass> {
void runOnOperation() override;
};
Expand All @@ -96,7 +98,4 @@ void ExtractAddressComputationGPUPass::runOnOperation() {
}
}

std::unique_ptr<Pass> createExtractAddressComputationGPUPass() {
return std::make_unique<ExtractAddressComputationGPUPass>();
}
} // namespace mlir::iree_compiler
Original file line number Diff line number Diff line change
Expand Up @@ -4,23 +4,26 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"

#define DEBUG_TYPE "iree-llvmgpu-cast-address-space-function"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_LLVMGPUCASTADDRESSSPACEFUNCTIONPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

namespace {

struct LLVMGPUCastAddressSpaceFunctionPass
: public LLVMGPUCastAddressSpaceFunctionBase<
struct LLVMGPUCastAddressSpaceFunctionPass final
: impl::LLVMGPUCastAddressSpaceFunctionPassBase<
LLVMGPUCastAddressSpaceFunctionPass> {
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<affine::AffineDialect, gpu::GPUDialect>();
Expand Down Expand Up @@ -75,10 +78,4 @@ struct LLVMGPUCastAddressSpaceFunctionPass
};

} // namespace

std::unique_ptr<OperationPass<ModuleOp>>
createLLVMGPUCastAddressSpaceFunction() {
return std::make_unique<LLVMGPUCastAddressSpaceFunctionPass>();
}

} // namespace mlir::iree_compiler
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h"
#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/VectorOpUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
Expand All @@ -21,6 +20,9 @@

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_LLVMGPUCASTTYPETOFITMMAPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

namespace {

struct UpcastContractOutput final : OpRewritePattern<vector::ContractionOp> {
Expand Down Expand Up @@ -72,9 +74,8 @@ struct UpcastContractOutput final : OpRewritePattern<vector::ContractionOp> {
}
};

struct LLVMGPUCastTypeToFitMMAPass
: public LLVMGPUCastTypeToFitMMABase<LLVMGPUCastTypeToFitMMAPass> {
public:
struct LLVMGPUCastTypeToFitMMAPass final
: impl::LLVMGPUCastTypeToFitMMAPassBase<LLVMGPUCastTypeToFitMMAPass> {
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<vector::VectorDialect>();
registry.insert<arith::ArithDialect>();
Expand Down Expand Up @@ -114,9 +115,4 @@ struct LLVMGPUCastTypeToFitMMAPass
}
};
} // namespace
std::unique_ptr<InterfacePass<FunctionOpInterface>>
createLLVMGPUCastTypeToFitMMAPass() {
return std::make_unique<LLVMGPUCastTypeToFitMMAPass>();
}

} // namespace mlir::iree_compiler
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
#include "iree/compiler/Codegen/LLVMGPU/PassDetail.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "llvm/ADT/SetVector.h"
Expand All @@ -27,6 +26,9 @@

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_LLVMGPUCONFIGUREVECTORLAYOUTSPASS
#include "iree/compiler/Codegen/LLVMGPU/Passes.h.inc"

namespace {

// Sets an anchoring layout for the given contraction op. Looks for a
Expand Down Expand Up @@ -279,10 +281,9 @@ LogicalResult setTransferReadAnchor(ArrayRef<int64_t> workgroupSize,
return success();
}

struct LLVMGPUConfigureVectorLayoutsPass
: public LLVMGPUConfigureVectorLayoutsBase<
struct LLVMGPUConfigureVectorLayoutsPass final
: impl::LLVMGPUConfigureVectorLayoutsPassBase<
LLVMGPUConfigureVectorLayoutsPass> {
public:
void getDependentDialects(DialectRegistry &registry) const override {
registry.insert<IREE::VectorExt::IREEVectorExtDialect>();
registry.insert<vector::VectorDialect>();
Expand Down Expand Up @@ -360,10 +361,4 @@ struct LLVMGPUConfigureVectorLayoutsPass
}
};
} // namespace

std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
createLLVMGPUConfigureVectorLayouts() {
return std::make_unique<LLVMGPUConfigureVectorLayoutsPass>();
}

} // namespace mlir::iree_compiler
Loading

0 comments on commit a72e78b

Please sign in to comment.