Bump to triton-lang/triton@142586f90aac

.github/workflows/test-plugin.yml

@@ -73,7 +73,7 @@ jobs:
       working-directory: triton
       run: |
         python3 -m pip install --upgrade pip
-        python3 -m pip install cmake==3.24 ninja pytest-xdist pybind11 setuptools
+        python3 -m pip install cmake==3.24 ninja pytest-xdist nanobind==2.10.2 setuptools
         sudo apt-get update -y
         sudo apt-get install -y ccache clang lld
         export TRITON_PLUGIN_DIRS="${GITHUB_WORKSPACE}/triton_shared"
@@ -154,7 +154,7 @@ jobs:
       working-directory: triton
       run: |
         python3 -m pip install --upgrade pip
-        python3 -m pip install cmake==3.24 ninja pytest-xdist pybind11 setuptools
+        python3 -m pip install cmake==3.24 ninja pytest-xdist nanobind==2.10.2 setuptools
         sudo apt-get update -y
         sudo apt-get install -y ccache clang lld
         export TRITON_PLUGIN_DIRS="${GITHUB_WORKSPACE}/triton_shared"

CMakeLists.txt

@@ -7,7 +7,6 @@ set(TRITON_SHARED_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}")
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) # Tablegen'd files
 include_directories(${Python3_INCLUDE_DIR})
-include_directories(${pybind11_INCLUDE_DIR})
 
 function(add_symlink source_file link_name)
     if (EXISTS ${source_file})
@@ -27,7 +26,7 @@ if (TRITON_SHARED_BUILD_CPU_BACKEND)
     get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
     add_triton_plugin(TritonShared ${CMAKE_CURRENT_SOURCE_DIR}/triton_shared.cc LINK_LIBS TritonSharedAnalysis TritonTilingExtIR ${conversion_libs} ${extension_libs} MLIRSparseTensorTransforms MLIRControlFlowTransforms MLIRTensorInferTypeOpInterfaceImpl)
-    target_link_libraries(TritonShared PRIVATE Python3::Module pybind11::headers)
+    target_link_libraries(TritonShared PRIVATE Python3::Module nanobind-static)
 endif()
 
 # Add symlinks to selected pytest files in triton. The tests are imported into triton-shared’s test folder to

backend/compiler.py

@@ -103,6 +103,7 @@ class CPUOptions:
     allowed_dot_input_precisions: Tuple[str] = ("ieee", )
     sanitize_overflow: bool = True
     instrumentation_mode: str = ""
+    fpsan_homomorphic_casts: bool = False
 
     def __post_init__(self):
         pass

include/triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtOps.td

@@ -166,6 +166,17 @@ class TritonTilingExt_TilingOp<string mnemonic> : Op<TritonTilingExt_Dialect, mn
       );
     }
 
+    FailureOr<TilingResult> $cppClass::getTiledImplementation(
+      OpBuilder &b,
+      ArrayRef<OpFoldResult> offsets,
+      ArrayRef<OpFoldResult> sizes,
+      ArrayRef<InnerTileAlignment> innerTileAlignments
+    ) {
+      return mlir::ttx::getTiledImplementation<$cppClass>(
+        *this, b, offsets, sizes
+      );
+    }
+
     // Forward each of the implementation to the shared implementation
     LogicalResult $cppClass::getResultTilePosition(
       OpBuilder &b,
@@ -192,6 +203,18 @@ class TritonTilingExt_TilingOp<string mnemonic> : Op<TritonTilingExt_Dialect, mn
       );
     }
 
+    FailureOr<TilingResult> $cppClass::generateResultTileValue(
+      OpBuilder &b,
+      unsigned resultNumber,
+      ArrayRef<OpFoldResult> offsets,
+      ArrayRef<OpFoldResult> sizes,
+      ArrayRef<InnerTileAlignment> innerTileAlignments
+    ) {
+      return mlir::ttx::generateResultTileValue<$cppClass>(
+        *this, b, resultNumber, offsets, sizes
+      );
+    }
+
     // Implemented as part of MemoryEffectsOpInterface
     void $cppClass::getEffects(
       SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>

lib/Transform/AddLLVMDebugInfo/AddLLVMDebugInfoPass.cpp

@@ -98,7 +98,7 @@ class AddLLVMDebugInfoPass
       LLVM::DISubprogramFlags subprogramFlags = setSubprogramFlags(funcOp);
 
       // retained nodes
-      llvm::ArrayRef<LLVM::DINodeAttr> importedModules;
+      llvm::ArrayRef<mlir::Attribute> importedModules;
 
       // annotations
       llvm::ArrayRef<LLVM::DINodeAttr> annotations;

python/examples/test_blockptr_complex_offset.py

@@ -9,31 +9,27 @@
 
 @triton.jit
 def block_copy_kernel(a_ptr, b_ptr):
-    a_block_ptr = tl.make_block_ptr(
+    a_desc = tl.make_tensor_descriptor(
         base=a_ptr + 8,
-        shape=(2, 2),
-        strides=(2, 1),
-        offsets=(0, 0),
-        block_shape=(2, 2),
-        order=(1, 0),
+        shape=[2, 4],
+        strides=[4, 1],
+        block_shape=[2, 4],
     )
-    b_block_ptr = tl.make_block_ptr(
+    a = a_desc.load([0, 0])
+    b_desc = tl.make_tensor_descriptor(
         base=b_ptr,
-        shape=(2, 2),
-        strides=(2, 1),
-        offsets=(0, 0),
-        block_shape=(2, 2),
-        order=(1, 0),
+        shape=[2, 4],
+        strides=[4, 1],
+        block_shape=[2, 4],
     )
-    a = tl.load(a_block_ptr, boundary_check=(0,))
-    tl.store(b_block_ptr, a, boundary_check=(0,))
+    b_desc.store([0, 0], a)
 
 
 
 def test(device):
-    input = torch.arange(0, 16, device=device, dtype=torch.float32)
-    output = torch.full((4,), -1, device=device, dtype=torch.float32)
-    expected = torch.arange(8, 12, device=device)
+    input = torch.arange(0, 32, device=device, dtype=torch.float32)
+    output = torch.full((8,), -1, device=device, dtype=torch.float32)
+    expected = torch.arange(8, 16, device=device)
     grid = lambda meta: (1,)
 
     block_copy_kernel[grid](input, output)

python/examples/test_load_2d_tensor_block.py

@@ -33,25 +33,21 @@ def kernel(
     pid0 = tl.program_id(axis=0)
     pid1 = tl.program_id(axis=1)
 
-    input_ptr = tl.make_block_ptr(
+    input_desc = tl.make_tensor_descriptor(
         base=x_ptr,
         shape=[n_rows, n_cols],
         strides=[stride_0, stride_1],
-        offsets=[pid0 * BLOCK_SIZE_ROW, pid1 * BLOCK_SIZE_COL],
         block_shape=[BLOCK_SIZE_ROW, BLOCK_SIZE_COL],
-        order=[1, 0],
     )
-    x = tl.load(input_ptr)
+    x = input_desc.load([pid0 * BLOCK_SIZE_ROW, pid1 * BLOCK_SIZE_COL])
     x = (2 * x) + 1
-    output_ptr = tl.make_block_ptr(
+    output_desc = tl.make_tensor_descriptor(
         base=y_ptr,
         shape=[n_rows, n_cols],
         strides=[stride_0, stride_1],
-        offsets=[pid0 * BLOCK_SIZE_ROW, pid1 * BLOCK_SIZE_COL],
         block_shape=[BLOCK_SIZE_ROW, BLOCK_SIZE_COL],
-        order=[1, 0],
     )
-    tl.store(output_ptr, x)
+    output_desc.store([pid0 * BLOCK_SIZE_ROW, pid1 * BLOCK_SIZE_COL], x)
 
 
 def test(device):
@@ -62,7 +58,7 @@ def test(device):
     )
     output = torch.full([n_rows, n_cols], -1, device=device, dtype=x.dtype)
     BLOCK_SIZE_ROW = 4
-    BLOCK_SIZE_COL = 2
+    BLOCK_SIZE_COL = 4
 
     grid = lambda meta: (n_rows // BLOCK_SIZE_ROW, n_cols // BLOCK_SIZE_COL)
 

python/examples/test_load_2d_tensor_col.py

@@ -15,7 +15,7 @@
 |     |     |     |     |
 |-----|-----|-----|-----|
 
-Each instance loads the entire column
+Each instance loads BLOCK_SIZE_COL columns
 """
 
 
@@ -29,37 +29,33 @@ def kernel(
     BLOCK_SIZE_COL: tl.constexpr,
 ):
     pid0 = tl.program_id(axis=0)
-    input_ptr = tl.make_block_ptr(
+    input_desc = tl.make_tensor_descriptor(
         base=x_ptr,
         shape=[n_rows, n_cols],
-        strides=[BLOCK_SIZE_COL, 1],
-        offsets=[0, pid0],
-        block_shape=[BLOCK_SIZE_ROW, 1],
-        order=[1, 0],
+        strides=[n_cols, 1],
+        block_shape=[BLOCK_SIZE_ROW, BLOCK_SIZE_COL],
     )
-    x = tl.load(input_ptr)
-    output_ptr = tl.make_block_ptr(
+    x = input_desc.load([0, pid0 * BLOCK_SIZE_COL])
+    output_desc = tl.make_tensor_descriptor(
         base=y_ptr,
         shape=[n_rows, n_cols],
-        strides=[BLOCK_SIZE_COL, 1],
-        offsets=[0, pid0],
-        block_shape=[BLOCK_SIZE_ROW, 1],
-        order=[1, 0],
+        strides=[n_cols, 1],
+        block_shape=[BLOCK_SIZE_ROW, BLOCK_SIZE_COL],
     )
-    tl.store(output_ptr, x)
+    output_desc.store([0, pid0 * BLOCK_SIZE_COL], x)
 
 
 def test(device):
     n_rows = 4
-    n_cols = 2
+    n_cols = 8
     x = torch.arange(0, n_rows * n_cols, 1, device=device, dtype=torch.float32).reshape(
         [n_rows, n_cols]
     )
     output = torch.full([n_rows, n_cols], -1, device=device, dtype=x.dtype)
     BLOCK_SIZE_ROW = n_rows
-    BLOCK_SIZE_COL = n_cols
+    BLOCK_SIZE_COL = 4
 
-    grid = lambda meta: (n_cols,)
+    grid = lambda meta: (n_cols // BLOCK_SIZE_COL,)
 
     kernel[grid](
         x,

python/examples/test_reduce.py

@@ -18,17 +18,13 @@ def reduce_kernel_2d(
     BLOCK_SIZE: tl.constexpr,
 ):
     pid0 = tl.program_id(axis=0)
-    x = tl.load(
-        tl.make_block_ptr(
-            base=x_ptr,
-            shape=[n_elements * tl.num_programs(0)],
-            strides=[1],
-            offsets=[stride * pid0],
-            block_shape=[BLOCK_SIZE],
-            order=[0],
-        ),
-        boundary_check=[0],
+    x_desc = tl.make_tensor_descriptor(
+        base=x_ptr,
+        shape=[n_elements * tl.num_programs(0)],
+        strides=[1],
+        block_shape=[BLOCK_SIZE],
     )
+    x = x_desc.load([stride * pid0])
     output = triton.language.sum(x, axis=0).to(dtype=x.dtype)
     tl.store(output_ptr + pid0, output)