[Backend][Hopper] Add a skeleton third-party warp specialization pass (#6624)

Adding a skeleton third-party warp specialization pass. It currently
doesn't do anything but serves as a placefolder for upcoming changes.

Author: htyu

bin/RegisterTritonDialects.h

@@ -17,6 +17,7 @@
 #include "triton/Dialect/TritonGPU/Transforms/Passes.h"
 #include "triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h"
 
+#include "nvidia/hopper/include/Transforms/Passes.h"
 #include "nvidia/include/Dialect/NVWS/Transforms/Passes.h"
 #include "nvidia/include/NVGPUToLLVM/Passes.h"
 #include "nvidia/include/TritonNVIDIAGPUToLLVM/Passes.h"
@@ -81,6 +82,9 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   // NVWS passes
   mlir::registerNVWSTransformsPasses();
 
+  // NVGPU transform passes
+  mlir::registerNVHopperTransformsPasses();
+
   registry.insert<
       mlir::triton::TritonDialect, mlir::cf::ControlFlowDialect,
       mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect,

third_party/nvidia/CMakeLists.txt

@@ -9,3 +9,4 @@ endif()
 if(TRITON_BUILD_UT)
   add_subdirectory(unittest)
 endif()
+add_subdirectory(hopper)

third_party/nvidia/hopper/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_subdirectory(include)
+add_subdirectory(lib)

third_party/nvidia/hopper/include/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(Transforms)

third_party/nvidia/hopper/include/Transforms/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name NVHopperTransforms)
+add_public_tablegen_target(NVHopperTransformsIncGen)

third_party/nvidia/hopper/include/Transforms/Passes.h

@@ -0,0 +1,18 @@
+
+#ifndef DIALECT_NV_TRANSFORMS_PASSES_H_
+#define DIALECT_NV_TRANSFORMS_PASSES_H_
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+// Generate the pass class declarations.
+#define GEN_PASS_DECL
+#include "nvidia/hopper/include/Transforms/Passes.h.inc"
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "nvidia/hopper/include/Transforms/Passes.h.inc"
+
+} // namespace mlir
+#endif // DIALECT_NV_TRANSFORMS_PASSES_H_

third_party/nvidia/hopper/include/Transforms/Passes.td

@@ -0,0 +1,23 @@
+#ifndef NV_TRANSFORMS_PASSES
+#define NV_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def NVGPUWarpSpecialization : Pass<"nvgpu-warp-specialization", "mlir::ModuleOp"> {
+  let summary = "Automaticl Warp specialization for NVIDIA GPU";
+
+  let description = [{
+    This pass automatically partitions user-defined kernels into
+    warp-specialized kernels, enabling finer-grained scheduling
+    and improved utilization of hardware resources.
+  }];
+
+  let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect"];
+  let options = [
+    Option<"numWarpGroups", "num-warp-groups",
+           "int32_t", /*default*/"0",
+           "number of warp groups for warp specialization">
+  ];
+}
+
+#endif // NV_TRANSFORMS_PASSES

third_party/nvidia/hopper/lib/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(Transforms)

third_party/nvidia/hopper/lib/Transforms/CMakeLists.txt

@@ -0,0 +1,11 @@
+add_triton_library(NVHopperTransforms
+  WarpSpecialization.cpp
+
+  DEPENDS
+  NVHopperTransformsIncGen
+
+  LINK_LIBS PUBLIC
+  TritonIR
+  TritonGPUIR
+  MLIRTransformUtils
+)

third_party/nvidia/hopper/lib/Transforms/WarpSpecialization.cpp

@@ -0,0 +1,32 @@
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+#include "nvidia/hopper/include/Transforms/Passes.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+
+#define DEBUG_TYPE "nvgpu-warp-specialization"
+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
+namespace mlir {
+
+#define GEN_PASS_DEF_NVGPUWARPSPECIALIZATION
+#include "nvidia/hopper/include/Transforms/Passes.h.inc"
+
+class NVGPUWarpSpecializationPass
+    : public impl::NVGPUWarpSpecializationBase<NVGPUWarpSpecializationPass> {
+public:
+  using impl::NVGPUWarpSpecializationBase<
+      NVGPUWarpSpecializationPass>::NVGPUWarpSpecializationBase;
+
+  void runOnFuncOp(triton::FuncOp funcOp) {
+    if (numWarpGroups <= 0)
+      return;
+  }
+
+  void runOnOperation() override {
+    getOperation()->walk([&](triton::FuncOp funcOp) { runOnFuncOp(funcOp); });
+  }
+};
+
+} // namespace mlir