diff --git a/mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt b/mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt --- a/mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt +++ b/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) diff --git a/mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt b/mlir/include/mlir/Dialect/NVGPU/Transforms/CMakeLists.txt copy from mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt copy to mlir/include/mlir/Dialect/NVGPU/Transforms/CMakeLists.txt --- a/mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/NVGPU/Transforms/CMakeLists.txt @@ -1,6 +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) diff --git a/mlir/include/mlir/Dialect/NVGPU/Passes.h b/mlir/include/mlir/Dialect/NVGPU/Transforms/Passes.h rename from mlir/include/mlir/Dialect/NVGPU/Passes.h rename to mlir/include/mlir/Dialect/NVGPU/Transforms/Passes.h --- a/mlir/include/mlir/Dialect/NVGPU/Passes.h +++ b/mlir/include/mlir/Dialect/NVGPU/Transforms/Passes.h @@ -18,7 +18,7 @@ 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 createOptimizeSharedMemoryPass(); @@ -31,7 +31,7 @@ /// 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 diff --git a/mlir/include/mlir/Dialect/NVGPU/Passes.td b/mlir/include/mlir/Dialect/NVGPU/Transforms/Passes.td rename from mlir/include/mlir/Dialect/NVGPU/Passes.td rename to mlir/include/mlir/Dialect/NVGPU/Transforms/Passes.td diff --git a/mlir/include/mlir/Dialect/NVGPU/Transforms/Utils.h b/mlir/include/mlir/Dialect/NVGPU/Transforms/Utils.h new file mode 100644 --- /dev/null +++ b/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 indices); + +} // namespace nvgpu +} // namespace mlir diff --git a/mlir/include/mlir/InitAllPasses.h b/mlir/include/mlir/InitAllPasses.h --- a/mlir/include/mlir/InitAllPasses.h +++ b/mlir/include/mlir/InitAllPasses.h @@ -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" diff --git a/mlir/lib/Dialect/NVGPU/Transforms/CMakeLists.txt b/mlir/lib/Dialect/NVGPU/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/NVGPU/Transforms/CMakeLists.txt +++ b/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 diff --git a/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp b/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp --- a/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp +++ b/mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp @@ -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" diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp --- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp +++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp @@ -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" @@ -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 @@ -107,38 +108,6 @@ permuteVectorOffset(builder, loc, indices, memrefTy, srcDim, tgtDim); } -Operation::operand_range getIndices(Operation *op) { - if (auto ldmatrixOp = dyn_cast(op)) - return ldmatrixOp.getIndices(); - if (auto copyOp = dyn_cast(op)) - return copyOp.getDstIndices(); - if (auto loadOp = dyn_cast(op)) - return loadOp.getIndices(); - if (auto storeOp = dyn_cast(op)) - return storeOp.getIndices(); - if (auto vectorReadOp = dyn_cast(op)) - return vectorReadOp.getIndices(); - if (auto vectorStoreOp = dyn_cast(op)) - return vectorStoreOp.getIndices(); - llvm_unreachable("unsupported op type"); -} - -void setIndices(Operation *op, ArrayRef indices) { - if (auto ldmatrixOp = dyn_cast(op)) - return ldmatrixOp.getIndicesMutable().assign(indices); - if (auto copyOp = dyn_cast(op)) - return copyOp.getDstIndicesMutable().assign(indices); - if (auto loadOp = dyn_cast(op)) - return loadOp.getIndicesMutable().assign(indices); - if (auto storeOp = dyn_cast(op)) - return storeOp.getIndicesMutable().assign(indices); - if (auto vectorReadOp = dyn_cast(op)) - return vectorReadOp.getIndicesMutable().assign(indices); - if (auto vectorStoreOp = dyn_cast(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 diff --git a/mlir/lib/Dialect/NVGPU/Transforms/Utils.cpp b/mlir/lib/Dialect/NVGPU/Transforms/Utils.cpp new file mode 100644 --- /dev/null +++ b/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(op)) + return ldmatrixOp.getIndices(); + if (auto copyOp = dyn_cast(op)) + return copyOp.getDstIndices(); + if (auto loadOp = dyn_cast(op)) + return loadOp.getIndices(); + if (auto storeOp = dyn_cast(op)) + return storeOp.getIndices(); + if (auto vectorReadOp = dyn_cast(op)) + return vectorReadOp.getIndices(); + if (auto vectorStoreOp = dyn_cast(op)) + return vectorStoreOp.getIndices(); + llvm_unreachable("unsupported op type"); +} + +void nvgpu::setIndices(Operation *op, ArrayRef indices) { + if (auto ldmatrixOp = dyn_cast(op)) + return ldmatrixOp.getIndicesMutable().assign(indices); + if (auto copyOp = dyn_cast(op)) + return copyOp.getDstIndicesMutable().assign(indices); + if (auto loadOp = dyn_cast(op)) + return loadOp.getIndicesMutable().assign(indices); + if (auto storeOp = dyn_cast(op)) + return storeOp.getIndicesMutable().assign(indices); + if (auto vectorReadOp = dyn_cast(op)) + return vectorReadOp.getIndicesMutable().assign(indices); + if (auto vectorStoreOp = dyn_cast(op)) + return vectorStoreOp.getIndicesMutable().assign(indices); + llvm_unreachable("unsupported op type"); +} diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -2752,11 +2752,11 @@ "-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"], ) @@ -2769,7 +2769,6 @@ ":GPUDialect", ":IR", ":NVGPUIncGen", - ":NVGPUPassIncGen", ":SideEffectInterfaces", "//llvm:Core", "//llvm:Support", @@ -2851,14 +2850,12 @@ 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",