MDRange iteration order transformation

mlir/include/lapis/Dialect/Kokkos/Transforms/Passes.h

@@ -31,6 +31,8 @@ std::unique_ptr<Pass> createKokkosMemorySpaceAssignmentPass();
 
 std::unique_ptr<Pass> createKokkosDualViewManagementPass();
 
+std::unique_ptr<Pass> createKokkosMdrangeIterationPass();
+
 //===----------------------------------------------------------------------===//
 // Registration.
 //===----------------------------------------------------------------------===//

mlir/include/lapis/Dialect/Kokkos/Transforms/Passes.td

@@ -56,4 +56,18 @@ def KokkosDualViewManagement : Pass<"kokkos-dualview-management", "ModuleOp"> {
   ];
 }
 
+def KokkosMdrangeIteration : Pass<"kokkos-mdrange-iteration", "ModuleOp"> {
+  let summary = "Rearange MDRange to improve memory access patterns on GPU";
+  let description = [{
+  }];
+  let constructor = "mlir::createKokkosMdrangeIterationPass()";
+  let dependentDialects = [
+    "arith::ArithDialect",
+    "func::FuncDialect",
+    "kokkos::KokkosDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect"
+  ];
+}
+
 #endif // MLIR_DIALECT_KOKKOS_TRANSFORMS_PASSES

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

@@ -4,6 +4,7 @@ add_mlir_dialect_library(MLIRKokkosTransforms
   KokkosLoopMapping.cpp
   KokkosMemorySpaceAssignment.cpp
   KokkosDualViewManagement.cpp
+  KokkosMdrangeIterationPass.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Kokkos