[Pipeliner] Merge warp specialization and pipeliner scheduling (#6887)

This PR refactors warp specialization to share the same scheduling as
software pipelining. What this means is that the pipeliner's loop
scheduler is used to set the stages and clusters of the ops, then on top
of that, warp specialization will perform partition assignment and split
the loop, introducing synchronization, into multiple loops that are then
individually software pipelined.

Author: Mogball

python/tutorials/06-fused-attention.py

@@ -15,7 +15,7 @@
 
 import pytest
 import torch
-import sys
+import os
 
 import triton
 import triton.language as tl
@@ -101,18 +101,18 @@ def _host_descriptor_pre_hook(nargs):
 if is_hip():
     NUM_STAGES_OPTIONS = [1]
 elif supports_host_descriptor():
-    NUM_STAGES_OPTIONS = [2, 3, 4, 5]
+    NUM_STAGES_OPTIONS = [2, 3, 4]
 else:
-    NUM_STAGES_OPTIONS = [2, 3, 4, 7]
+    NUM_STAGES_OPTIONS = [2, 3, 4]
 
 configs = [
     triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w, pre_hook=_host_descriptor_pre_hook) \
-    for BM in [64, 128, 256]\
+    for BM in [64, 128]\
     for BN in [64, 128]\
     for s in NUM_STAGES_OPTIONS \
     for w in [4, 8]\
 ]
-if "pytest" in sys.modules:
+if "PYTEST_VERSION" in os.environ:
     # Use a single config in testing for reproducibility
     configs = [
         triton.Config(dict(BLOCK_M=64, BLOCK_N=64), num_stages=4, num_warps=4, pre_hook=_host_descriptor_pre_hook),
@@ -503,6 +503,8 @@ def grid(META):
             return (triton.cdiv(q.shape[2], META["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
 
         ctx.grid = grid
+        if is_cuda() and warp_specialize:
+            extra_kern_args["maxnreg"] = 80
         _attn_fwd[grid](
             sm_scale, M,  #
             q.shape[0], q.shape[1],  #