[mlir][NVGPU][NFC] Clean up code structure

[springerm]

Summary

• Move passes to Transforms directory.
• Add Utils.h (will be utilized in a subsequent change).

mlir/include/mlir/Dialect/NVGPU/CMakeLists.txt

 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)

mlir/include/mlir/Dialect/NVGPU/Transforms/CMakeLists.txt

-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)

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

 #define MLIR_DIALECT_NVGPU_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 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();
 
 } // namespace nvgpu
 
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// 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
 
 #endif // MLIR_DIALECT_NVGPU_PASSES_H_

mlir/include/mlir/Dialect/NVGPU/Transforms/Passes.td

 //===-- Passes.td - NvGpu pass definition file ------------*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_NVGPU_PASSES_TD_
 #define MLIR_DIALECT_NVGPU_PASSES_TD_
 
 include "mlir/Pass/PassBase.td"
 
 def OptimizeSharedMemory : Pass<"nvgpu-optimize-shared-memory"> {
   let summary = "Optimizes accesses to shard memory memrefs in order to reduce bank conflicts.";
   let constructor = "mlir::nvgpu::createOptimizeSharedMemoryPass()";
   let dependentDialects = [
     "memref::MemRefDialect", "vector::VectorDialect"
   ];
 }
 
 #endif // MLIR_DIALECT_NVGPU_PASSES_TD_

mlir/include/mlir/Dialect/NVGPU/Transforms/Utils.h

+//===- 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

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
 #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"
 #include "mlir/Dialect/SparseTensor/Pipelines/Passes.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
 #include "mlir/Dialect/Transform/Transforms/Passes.h"

mlir/lib/Dialect/NVGPU/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRNVGPUTransforms
   OptimizeSharedMemory.cpp
-  MmaSyncTF32Transform.cpp  
+  MmaSyncTF32Transform.cpp
+  Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/NVGPU
 
   DEPENDS
   MLIRNVGPUPassIncGen
 
   LINK_LIBS PUBLIC

mlir/lib/Dialect/NVGPU/Transforms/MmaSyncTF32Transform.cpp

 //===- OptimizeSharedMemory.cpp - MLIR NVGPU pass implementation ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements transforms to enable 1xtf32 and 3xtf32 nvgpu.mma sync
 // operations on f32 input datatype
 //
 //===----------------------------------------------------------------------===//
 
+#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"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/MathExtras.h"
 
 using namespace mlir;
 using namespace mlir::nvgpu;

mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp

 //===- OptimizeSharedMemory.cpp - MLIR NVGPU pass implementation ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements transforms to optimize accesses to shared memory.
 //
 //===----------------------------------------------------------------------===//
 
-#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"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/MathExtras.h"
 
 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
 
 using namespace mlir;
 using namespace mlir::nvgpu;
 
 /// The size of a shared memory line according to NV documentation.
 constexpr int64_t kSharedMemoryLineSizeBytes = 128;

 static void transformIndices(OpBuilder &builder, Location loc,
                              SmallVector<Value, 4> &indices,
                              MemRefType memrefTy, int64_t srcDim,
                              int64_t tgtDim) {
   indices[tgtDim] =
       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
 getShmReadAndWriteOps(Operation *parentOp, Value shmMemRef,
                       SmallVector<Operation *, 16> &readOps,
                       SmallVector<Operation *, 16> &writeOps) {
   parentOp->walk([&](Operation *op) {
     MemoryEffectOpInterface iface = dyn_cast<MemoryEffectOpInterface>(op);

mlir/lib/Dialect/NVGPU/Transforms/Utils.cpp

+//===- 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");
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

gentbl_cc_library(
     name = "NVGPUPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             [
                 "-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"],
 )
 
 cc_library(
     name = "NVGPUDialect",
     srcs = ["lib/Dialect/NVGPU/IR/NVGPUDialect.cpp"],
     hdrs = ["include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h"],
     includes = ["include"],
     deps = [
         ":GPUDialect",
         ":IR",
         ":NVGPUIncGen",
-        ":NVGPUPassIncGen",
         ":SideEffectInterfaces",
         "//llvm:Core",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "NVGPUTransformOps",
deps = [
         ":NVGPUDialect",
         ":NVVMDialect",
         ":VectorDialect",
     ],
 )
 
 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",
         ":ArithDialect",
         ":FuncDialect",
         ":GPUDialect",
         ":IR",
         ":MemRefDialect",