Merge OpenAI Triton commit 27c5940 (#4940)

This PR change the Triton base from
05b2c18 to
27c5940 (Aug 11).
Pass rate: 98.85%

Author: whitneywhtsang

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -48,6 +48,9 @@ class TargetInfoBase {
   virtual Value shuffleIdx(RewriterBase &rewriter, Location loc, Value val,
                            Value i) const = 0;
 
+  virtual Value permute(RewriterBase &rewriter, Location loc, Value a, Value b,
+                        Value selector) const = 0;
+
   virtual Value programId(RewriterBase &rewriter, Location loc,
                           ModuleOp moduleOp, ProgramIDDim axis) const = 0;
 

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -571,8 +571,9 @@ SmallVector<Value> lowerLdSt(
     ArrayRef<Value> valsArray, // Input for store, output for load
     Type llvmElemTy, Value smemBase,
     std::function<Value(Value)> calcPaddedOffset, Value affineOffset,
-    uint64_t maskSpanAffineOffset, RewriterBase &rewriter,
-    const TargetInfoBase &targetInfo, std::optional<int> maybeMaxVecElems,
+    uint64_t maskSpanAffineOffset, Value laneId, Value warpId,
+    RewriterBase &rewriter, const TargetInfoBase &targetInfo,
+    std::optional<int> maybeMaxVecElems,
     std::function<SmallVector<Value>(RewriterBase &, Location, ArrayRef<Value>,
                                      Value, int, VectorType)>
         lowerInst);

lib/Conversion/TritonGPUToLLVM/Utility.cpp

@@ -542,18 +542,20 @@ lowerLdStShared(Location loc, MLIRContext *ctx, LinearLayout cvt,
       return unpackLLVector(loc, valsVec, rewriter);
     }
   };
+  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
   return lowerLdSt(loc, ctx, cvt, valsArray, llvmElemTy, smemBase,
-                   calcPaddedOffset, affineOffset, maskSpanAffineOffset,
-                   rewriter, targetInfo, {}, emitLdSt);
+                   calcPaddedOffset, affineOffset, maskSpanAffineOffset, laneId,
+                   warpId, rewriter, targetInfo, {}, emitLdSt);
 }
 
 SmallVector<Value> lowerLdSt(
     Location loc, MLIRContext *ctx, LinearLayout cvt,
     ArrayRef<Value> valsArray, // Input for store, output for load
     Type llvmElemTy, Value smemBase,
     std::function<Value(Value)> calcPaddedOffset, Value affineOffset,
-    uint64_t maskSpanAffineOffset, RewriterBase &rewriter,
-    const TargetInfoBase &targetInfo, std::optional<int> maybeMaxVecElems,
+    uint64_t maskSpanAffineOffset, Value laneId, Value warpId,
+    RewriterBase &rewriter, const TargetInfoBase &targetInfo,
+    std::optional<int> maybeMaxVecElems,
     std::function<SmallVector<Value>(RewriterBase &, Location, ArrayRef<Value>,
                                      Value, int, VectorType)>
         lowerInst) {
@@ -599,7 +601,6 @@ SmallVector<Value> lowerLdSt(
       zerosLike(LinearLayout::identity1D(bitwidth / 8, kReg, kOffset));
   auto i8AddrLayout = i8Tile * addrLayout;
 
-  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
   auto regBaseI8 =
       applyLinearLayout(
           loc, rewriter, i8AddrLayout,
@@ -2022,16 +2023,17 @@ void finalizeTensorAtomicResults(Operation *op, RankedTensorType tensorTy,
   };
 
   auto noPaddingOffset = [](Value v) { return v; };
+  auto [laneId, warpId] = getLaneAndWarpId(rewriter, loc);
   lowerLdSt(loc, ctx, dstLayout, resultVals, valueElemTy, smemBase,
             /*calcPaddedOffset=*/noPaddingOffset, /*affineOffset=*/b.i32_val(0),
-            /*maskSpanAffineOffset=*/0, rewriter, targetInfo,
+            /*maskSpanAffineOffset=*/0, laneId, warpId, rewriter, targetInfo,
             /*maybeMaxVecElems=*/{}, emitSt);
   b.barrier();
   resultVals = lowerLdSt(loc, ctx, dstLayout, resultVals, valueElemTy, smemBase,
                          /*calcPaddedOffset=*/noPaddingOffset,
                          /*affineOffset=*/b.i32_val(0),
-                         /*maskSpanAffineOffset=*/0, rewriter, targetInfo,
-                         /*maybeMaxVecElems=*/{}, emitLd);
+                         /*maskSpanAffineOffset=*/0, laneId, warpId, rewriter,
+                         targetInfo, /*maybeMaxVecElems=*/{}, emitLd);
 
   // Create the result struct and replace the operation
   Value resultStruct =

python/test/conftest.py

@@ -108,24 +108,24 @@ def fresh_cache():
 
 
 @pytest.fixture
-def fresh_knobs(monkeypatch):
+def fresh_knobs():
     from triton._internal_testing import _fresh_knobs_impl
-    fresh_function, reset_function = _fresh_knobs_impl(monkeypatch)
+    fresh_function, reset_function = _fresh_knobs_impl()
     try:
         yield fresh_function()
     finally:
         reset_function()
 
 
 @pytest.fixture
-def fresh_knobs_except_libraries(monkeypatch):
+def fresh_knobs_except_libraries():
     """
     A variant of `fresh_knobs` that keeps library path
     information from the environment as these may be
     needed to successfully compile kernels.
     """
     from triton._internal_testing import _fresh_knobs_impl
-    fresh_function, reset_function = _fresh_knobs_impl(monkeypatch, skipped_attr={"build", "nvidia", "amd"})
+    fresh_function, reset_function = _fresh_knobs_impl(skipped_attr={"build", "nvidia", "amd"})
     try:
         yield fresh_function()
     finally:

python/test/unit/cuda/test_tma_descriptor.py

@@ -55,9 +55,12 @@ def example_load_store_kernel(X, Y, x_off, y_off, x_size, y_size):
     store_ragged(Y, y_off, y_size, [0, 0], data)
 
 
-@pytest.mark.parametrize("write_only", [False, True])
-@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
-def test_ragged_tma(dtype, write_only):
+@pytest.mark.parametrize("dtype", [
+    "bfloat16", "float16", "float32", "float64",  # floating-point
+    "int8", "int16", "int32", "int64",  # signed integers
+    "uint8", "uint16", "uint32", "uint64"  # unsigned integers
+])
+def test_ragged_tma(dtype):
 
     if not torch.cuda.is_available() or not torch.cuda.get_device_capability()[0] >= 9:
         pytest.skip("Test requires Hopper or Blackwell target.")
@@ -67,10 +70,10 @@ def test_ragged_tma(dtype, write_only):
 
     src = torch.randn((1024, 80), dtype=torch.float32, device="cuda").to(dtype)
     ref = torch.randn((1024, 80), dtype=torch.float32, device="cuda").to(dtype)
-    dst = 1.0 * ref
+    dst = ref.clone()
 
     X = create_ragged_descriptor(src, [32, 128])
-    Y = create_ragged_descriptor(dst, [32, 128], write_only=write_only)
+    Y = create_ragged_descriptor(dst, [32, 128])
 
     x_off = 42
     y_off = 51

python/test/unit/language/test_core.py

@@ -7325,6 +7325,26 @@ def simple(data, out):
         assert amdgcn_gfx.group(1) == arch
 
 
+def test_num_ctas_pre_sm90(device):
+    if not is_cuda() and not is_hip():
+        pytest.skip("Only supported on CUDA and HIP")
+
+    @triton.jit
+    def _kernel(src):
+        pass
+
+    src = torch.empty(1, device=device)
+    if is_cuda():
+        arch = "sm80"
+        msg = r"num_ctas > 1 requires NVIDIA SM90\+ \(Hopper\)"
+    else:
+        arch = "gfx942"
+        msg = r"num_ctas > 1 not supported for AMD GPUs"
+
+    with pytest.raises(ValueError, match=msg):
+        _kernel.warmup(src, grid=(1, ), num_ctas=2, arch=arch)
+
+
 # -----------------------
 # test propagate_nan
 # -----------------------

python/test/unit/runtime/test_driver.py

@@ -10,11 +10,11 @@ def test_is_lazy():
     from importlib import reload
     reload(sys.modules["triton.runtime.driver"])
     reload(sys.modules["triton.runtime"])
-    mod = sys.modules[triton.runtime.driver.__module__]
-    assert isinstance(triton.runtime.driver.active, getattr(mod, "LazyProxy"))
-    assert triton.runtime.driver.active._obj is None
+    assert triton.runtime.driver._active is None
+    assert triton.runtime.driver._default is None
+    assert isinstance(triton.runtime.driver.active, getattr(triton.backends.driver, "DriverBase"))
+    assert isinstance(triton.runtime.driver.default, getattr(triton.backends.driver, "DriverBase"))
     utils = triton.runtime.driver.active.utils  # noqa: F841
-    assert issubclass(triton.runtime.driver.active._obj.__class__, getattr(triton.backends.driver, "DriverBase"))
 
 
 def test_kernel_in_thread(device):

python/triton/_internal_testing.py

@@ -204,12 +204,14 @@ def unwrap_tensor(t: Union[torch.Tensor, triton.runtime.jit.TensorWrapper]) -> t
     return t
 
 
-def _fresh_knobs_impl(monkeypatch, skipped_attr: Optional[Set[str]] = None):
+def _fresh_knobs_impl(skipped_attr: Optional[Set[str]] = None):
     from triton import knobs
 
     if skipped_attr is None:
         skipped_attr = set()
 
+    monkeypatch = pytest.MonkeyPatch()
+
     knobs_map = {
         name: knobset
         for name, knobset in knobs.__dict__.items()
@@ -237,6 +239,9 @@ def fresh_function():
     def reset_function():
         for name, knobset in knobs_map.items():
             setattr(knobs, name, knobset)
+        # `undo` should be placed before `del os.environ`
+        # Otherwise, it may restore environment variables that monkeypatch deleted
+        monkeypatch.undo()
         for k in env_to_unset:
             if k in os.environ:
                 del os.environ[k]

python/triton/compiler/compiler.py

@@ -442,13 +442,14 @@ def __init__(self, src, metadata_group, hash):
         # (e.g., checking amount of shared memory on current device)
         self.module = None
         self.function = None
+        self._run = None
 
     def _init_handles(self):
         if self.module is not None:
             return
         device = driver.active.get_current_device()
         # create launcher
-        self.run = driver.active.launcher_cls(self.src, self.metadata)
+        self._run = driver.active.launcher_cls(self.src, self.metadata)
         # not enough shared memory to run the kernel
         max_shared = max_shared_mem(device)
         if self.metadata.shared > max_shared:
@@ -469,10 +470,14 @@ def _init_handles(self):
         if self.metadata.num_warps * warp_size > self.n_max_threads:
             raise OutOfResources(self.metadata.num_warps * warp_size, self.n_max_threads, "threads")
 
-    def __getattribute__(self, name):
-        if name == 'run':
+    @property
+    def run(self):
+        # it should be safe to do this as launch_metadata will
+        # call _init_handles before running the kernel or it
+        # was called manually or it was already initialized
+        if self._run is None:
             self._init_handles()
-        return super().__getattribute__(name)
+        return self._run
 
     def launch_metadata(self, grid, stream, *args):
         if knobs.runtime.launch_enter_hook is None:

python/triton/runtime/driver.py

@@ -2,8 +2,6 @@
 
 from ..backends import backends, DriverBase
 
-from typing import Any, Callable, Generic, TypeVar, Union
-
 
 def _create_driver() -> DriverBase:
     active_drivers = [x.driver for x in backends.values() if x.driver.is_active()]
@@ -12,52 +10,29 @@ def _create_driver() -> DriverBase:
     return active_drivers[0]()
 
 
-T = TypeVar("T")
-
-
-class LazyProxy(Generic[T]):
-
-    def __init__(self, init_fn: Callable[[], T]) -> None:
-        self._init_fn = init_fn
-        self._obj: Union[T, None] = None
-
-    def _initialize_obj(self) -> T:
-        if self._obj is None:
-            self._obj = self._init_fn()
-        return self._obj
-
-    def __getattr__(self, name) -> Any:
-        return getattr(self._initialize_obj(), name)
-
-    def __setattr__(self, name: str, value: Any) -> None:
-        if name in ["_init_fn", "_obj"]:
-            super().__setattr__(name, value)
-        else:
-            setattr(self._initialize_obj(), name, value)
-
-    def __delattr__(self, name: str) -> None:
-        delattr(self._initialize_obj(), name)
-
-    def __repr__(self) -> str:
-        if self._obj is None:
-            return f"<{self.__class__.__name__} for {self._init_fn} not yet initialized>"
-        return repr(self._obj)
-
-    def __str__(self) -> str:
-        return str(self._initialize_obj())
-
-
 class DriverConfig:
 
     def __init__(self) -> None:
-        self.default: LazyProxy[DriverBase] = LazyProxy(_create_driver)
-        self.active: Union[LazyProxy[DriverBase], DriverBase] = self.default
+        self._default: DriverBase | None = None
+        self._active: DriverBase | None = None
+
+    @property
+    def default(self) -> DriverBase:
+        if self._default is None:
+            self._default = _create_driver()
+        return self._default
+
+    @property
+    def active(self) -> DriverBase:
+        if self._active is None:
+            self._active = self.default
+        return self._active
 
     def set_active(self, driver: DriverBase) -> None:
-        self.active = driver
+        self._active = driver
 
     def reset_active(self) -> None:
-        self.active = self.default
+        self._active = self.default
 
 
 driver = DriverConfig()