Skip to content

Commit

Permalink
[mlir][NVGPU][NFC] Clean up code structure
Browse files Browse the repository at this point in the history
* Move passes to `Transforms` directory.
* Add `Utils.h` (will be utilized in a subsequent change).

Differential Revision: https://reviews.llvm.org/D155427
  • Loading branch information
matthias-springer committed Jul 17, 2023
1 parent 92542f2 commit a4f4d82
Show file tree
Hide file tree
Showing 11 changed files with 95 additions and 59 deletions.
9 changes: 1 addition & 8 deletions mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt
@@ -1,10 +1,3 @@
add_subdirectory(IR)
add_subdirectory(TransformOps)

set(LLVM_TARGET_DEFINITIONS Passes.td)
mlir_tablegen(Passes.h.inc -gen-pass-decls -name NVGPU)
mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix NVGPU)
mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix NVGPU)
add_public_tablegen_target(MLIRNVGPUPassIncGen)

add_mlir_doc(Passes NVGPUPasses ./ -gen-pass-doc)
add_subdirectory(Transforms)
7 changes: 7 additions & 0 deletions mlir/include/mlir/Dialect/NVGPU/Transforms/CMakeLists.txt
@@ -0,0 +1,7 @@
set(LLVM_TARGET_DEFINITIONS Passes.td)
mlir_tablegen(Passes.h.inc -gen-pass-decls -name NVGPU)
mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix NVGPU)
mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix NVGPU)
add_public_tablegen_target(MLIRNVGPUPassIncGen)

add_mlir_doc(Passes NVGPUPasses ./ -gen-pass-doc)
Expand Up @@ -18,7 +18,7 @@ namespace mlir {
namespace nvgpu {

#define GEN_PASS_DECL
#include "mlir/Dialect/NVGPU/Passes.h.inc"
#include "mlir/Dialect/NVGPU/Transforms/Passes.h.inc"

/// Create a pass to optimize shared memory reads and writes.
std::unique_ptr<Pass> createOptimizeSharedMemoryPass();
Expand All @@ -31,7 +31,7 @@ std::unique_ptr<Pass> createOptimizeSharedMemoryPass();

/// Generate the code for registering passes.
#define GEN_PASS_REGISTRATION
#include "mlir/Dialect/NVGPU/Passes.h.inc"
#include "mlir/Dialect/NVGPU/Transforms/Passes.h.inc"

} // namespace mlir

Expand Down
File renamed without changes.
21 changes: 21 additions & 0 deletions mlir/include/mlir/Dialect/NVGPU/Transforms/Utils.h
@@ -0,0 +1,21 @@
//===- Utils.h - Transform utilities -----------------------------*- 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/IR/Operation.h"

namespace mlir {
namespace nvgpu {

/// Get the indices that the given load/store operation is operating on.
Operation::operand_range getIndices(Operation *op);

/// Set the indices that the given load/store operation is operating on.
void setIndices(Operation *op, ArrayRef<Value> indices);

} // namespace nvgpu
} // namespace mlir
2 changes: 1 addition & 1 deletion mlir/include/mlir/InitAllPasses.h
Expand Up @@ -27,7 +27,7 @@
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/NVGPU/Passes.h"
#include "mlir/Dialect/NVGPU/Transforms/Passes.h"
#include "mlir/Dialect/SCF/Transforms/Passes.h"
#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
#include "mlir/Dialect/Shape/Transforms/Passes.h"
Expand Down
3 changes: 2 additions & 1 deletion mlir/lib/Dialect/NVGPU/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
add_mlir_dialect_library(MLIRNVGPUTransforms
OptimizeSharedMemory.cpp
MmaSyncTF32Transform.cpp
MmaSyncTF32Transform.cpp
Utils.cpp

ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/NVGPU
Expand Down
4 changes: 2 additions & 2 deletions mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp
Expand Up @@ -11,12 +11,12 @@
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/NVGPU/Transforms/Transforms.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/NVGPU/Passes.h"
#include "mlir/Dialect/NVGPU/Transforms/Transforms.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LogicalResult.h"
Expand Down
37 changes: 3 additions & 34 deletions mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
Expand Up @@ -10,13 +10,14 @@
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/NVGPU/Passes.h"
#include "mlir/Dialect/NVGPU/Transforms/Passes.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/NVGPU/Transforms/Transforms.h"
#include "mlir/Dialect/NVGPU/Transforms/Utils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LogicalResult.h"
Expand All @@ -26,7 +27,7 @@
namespace mlir {
namespace nvgpu {
#define GEN_PASS_DEF_OPTIMIZESHAREDMEMORY
#include "mlir/Dialect/NVGPU/Passes.h.inc"
#include "mlir/Dialect/NVGPU/Transforms/Passes.h.inc"
} // namespace nvgpu
} // namespace mlir

Expand Down Expand Up @@ -107,38 +108,6 @@ static void transformIndices(OpBuilder &builder, Location loc,
permuteVectorOffset(builder, loc, indices, memrefTy, srcDim, tgtDim);
}

Operation::operand_range getIndices(Operation *op) {
if (auto ldmatrixOp = dyn_cast<LdMatrixOp>(op))
return ldmatrixOp.getIndices();
if (auto copyOp = dyn_cast<DeviceAsyncCopyOp>(op))
return copyOp.getDstIndices();
if (auto loadOp = dyn_cast<memref::LoadOp>(op))
return loadOp.getIndices();
if (auto storeOp = dyn_cast<memref::StoreOp>(op))
return storeOp.getIndices();
if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))
return vectorReadOp.getIndices();
if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))
return vectorStoreOp.getIndices();
llvm_unreachable("unsupported op type");
}

void setIndices(Operation *op, ArrayRef<Value> indices) {
if (auto ldmatrixOp = dyn_cast<LdMatrixOp>(op))
return ldmatrixOp.getIndicesMutable().assign(indices);
if (auto copyOp = dyn_cast<DeviceAsyncCopyOp>(op))
return copyOp.getDstIndicesMutable().assign(indices);
if (auto loadOp = dyn_cast<memref::LoadOp>(op))
return loadOp.getIndicesMutable().assign(indices);
if (auto storeOp = dyn_cast<memref::StoreOp>(op))
return storeOp.getIndicesMutable().assign(indices);
if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))
return vectorReadOp.getIndicesMutable().assign(indices);
if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))
return vectorStoreOp.getIndicesMutable().assign(indices);
llvm_unreachable("unsupported op type");
}

/// Return all operations within `parentOp` that read from or write to
/// `shmMemRef`.
static LogicalResult
Expand Down
48 changes: 48 additions & 0 deletions mlir/lib/Dialect/NVGPU/Transforms/Utils.cpp
@@ -0,0 +1,48 @@
//===- Utils.cpp - Transform utilities ------------------------------------===//
//
// 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/NVGPU/Transforms/Utils.h"

#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"

using namespace mlir;
using namespace mlir::nvgpu;

Operation::operand_range nvgpu::getIndices(Operation *op) {
if (auto ldmatrixOp = dyn_cast<LdMatrixOp>(op))
return ldmatrixOp.getIndices();
if (auto copyOp = dyn_cast<DeviceAsyncCopyOp>(op))
return copyOp.getDstIndices();
if (auto loadOp = dyn_cast<memref::LoadOp>(op))
return loadOp.getIndices();
if (auto storeOp = dyn_cast<memref::StoreOp>(op))
return storeOp.getIndices();
if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))
return vectorReadOp.getIndices();
if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))
return vectorStoreOp.getIndices();
llvm_unreachable("unsupported op type");
}

void nvgpu::setIndices(Operation *op, ArrayRef<Value> indices) {
if (auto ldmatrixOp = dyn_cast<LdMatrixOp>(op))
return ldmatrixOp.getIndicesMutable().assign(indices);
if (auto copyOp = dyn_cast<DeviceAsyncCopyOp>(op))
return copyOp.getDstIndicesMutable().assign(indices);
if (auto loadOp = dyn_cast<memref::LoadOp>(op))
return loadOp.getIndicesMutable().assign(indices);
if (auto storeOp = dyn_cast<memref::StoreOp>(op))
return storeOp.getIndicesMutable().assign(indices);
if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))
return vectorReadOp.getIndicesMutable().assign(indices);
if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))
return vectorStoreOp.getIndicesMutable().assign(indices);
llvm_unreachable("unsupported op type");
}
19 changes: 8 additions & 11 deletions utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
Expand Up @@ -2752,11 +2752,11 @@ gentbl_cc_library(
"-gen-pass-decls",
"-name=NVGPU",
],
"include/mlir/Dialect/NVGPU/Passes.h.inc",
"include/mlir/Dialect/NVGPU/Transforms/Passes.h.inc",
),
],
tblgen = ":mlir-tblgen",
td_file = "include/mlir/Dialect/NVGPU/Passes.td",
td_file = "include/mlir/Dialect/NVGPU/Transforms/Passes.td",
deps = [":PassBaseTdFiles"],
)

Expand All @@ -2769,7 +2769,6 @@ cc_library(
":GPUDialect",
":IR",
":NVGPUIncGen",
":NVGPUPassIncGen",
":SideEffectInterfaces",
"//llvm:Core",
"//llvm:Support",
Expand Down Expand Up @@ -2851,14 +2850,12 @@ cc_library(

cc_library(
name = "NVGPUTransforms",
srcs = [
"lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp",
"lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp",
],
hdrs = [
"include/mlir/Dialect/NVGPU/Passes.h",
"include/mlir/Dialect/NVGPU/Transforms/Transforms.h",
],
srcs = glob([
"lib/Dialect/NVGPU/Transforms/*.cpp",
]),
hdrs = glob([
"include/mlir/Dialect/NVGPU/Transforms/*.h",
]),
includes = ["include"],
deps = [
":AffineDialect",
Expand Down

0 comments on commit a4f4d82

Please sign in to comment.