[intel] 2Dblock runtime HW checks

Author: januszjah

include/triton/Tools/Sys/GetEnv.hpp

@@ -44,6 +44,7 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_F32_DEFAULT",
     "TRITON_PREFER_TMEM_16x256_LAYOUT",
     "TRITON_ENABLE_EXPERIMENTAL_CONSAN",
+    "TRITON_INTEL_2DBLOCK_ASSERT",
     "TRITON_INTEL_AGGRESSIVE_DPAS_REUSE",
     "TRITON_INTEL_ENABLE_BLOCK_IO_ALL_LAYOUTS",
     "TRITON_INTEL_ENABLE_DPAS_FOR_WARP_SIZE_32",

python/test/unit/intel/block_io_helper.py

@@ -0,0 +1,54 @@
+import torch
+import triton
+
+import ctypes
+import sys
+
+
+def run_ir(device, temp_file):
+    ir = r"""
+    module attributes {
+        ttg.target = "xpu",
+        "ttg.num-warps" = 32 : i32,
+        "ttg.num-ctas" = 1 : i32,
+        "ttg.threads-per-warp" = 16 : i32
+        } {
+        tt.func @dyn_block(
+            %iptr : i64, %base_width : i32,
+            %base_height : i32, %base_pitch : i32,
+            %x : i32, %y : i32) {
+            %p0 = llvm.inttoptr %iptr : i64 to !llvm.ptr
+
+            %v = triton_gen.2Dblockload %p0, %base_width, %base_height,
+                %base_pitch, %x, %y
+                { elem_size_in_bits = 8, tile_width = 8, tile_height = 8,
+                v_blocks = 1, transpose = false,
+                vnni_transform = false, cache_control = Default }
+                : (!llvm.ptr, i32, i32, i32, i32, i32)
+                -> vector<2xi16>
+
+            // prevent GluonInline
+            %v_i32 = llvm.bitcast %v : vector<2xi16> to i32
+            llvm.inline_asm has_side_effects asm_dialect = att
+                "", "r" %v_i32 : (i32) -> ()
+
+            tt.return
+        }
+    }
+    """
+
+    with open(temp_file, "w", encoding="utf-8") as f:
+        f.write(ir)
+
+    kernel = triton.compile(temp_file)
+
+    a = torch.randn((256, 64), dtype=torch.float32, device=device)
+
+    addr = ctypes.c_int64(a.data_ptr()).value
+
+    kernel[(1, 1, 1)](addr, 64, 64, 1, 0, 0)
+
+
+if __name__ == "__main__":
+    fn = globals()[sys.argv[1]]
+    fn(*sys.argv[2:])

python/test/unit/intel/test_block_load.py

@@ -207,3 +207,30 @@ def triton_mm(X, Y, b=None, transpose_x=False, transpose_y=False):
     result_tor = fn_tor()
     result_tri = fn_tri()
     torch.testing.assert_close(result_tri, result_tor, atol=1e-2, rtol=1e-3)
+
+
+@pytest.mark.skipif(not is_xpu(), reason="Block load tests are specific to the XPU backend")
+@pytest.mark.xfail(
+    not (torch.xpu.get_device_capability()['has_subgroup_2d_block_io']
+         and torch.xpu.get_device_capability()['has_subgroup_matrix_multiply_accumulate']),
+    reason="Block loads and/or DPAS not supported on this architecture", run=False)
+def test_block_load_asserts(monkeypatch, device, tmp_path: pathlib.Path):
+    monkeypatch.setenv("TRITON_INTEL_2DBLOCK_ASSERT", "1")
+
+    import os
+    import signal
+    import subprocess
+    import sys
+
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    helper_path = os.path.join(dir_path, "block_io_helper.py")
+
+    temp_file = tmp_path / "test_block_load_asserts.ttgir"
+
+    proc = subprocess.run(
+        [sys.executable, helper_path, "run_ir", device, str(temp_file)],
+        capture_output=True,
+    )
+
+    rc = proc.returncode
+    assert rc == -signal.SIGABRT or rc == (128 + signal.SIGABRT)

test/TritonGEN/tritongen-2Dblockload-to-llvm-asserts.mlir

@@ -0,0 +1,27 @@
+// RUN: env TRITON_INTEL_2DBLOCK_ASSERT=1 triton-opt -convert-tritongen-to-llvm -split-input-file %s | FileCheck %s
+
+module attributes {"ttg.threads-per-warp" = 16 : i32} {
+llvm.func @triton_gen.2Dblockload(%ptr : !llvm.ptr<1>, %base_width : i32, %base_height : i32, %base_pitch : i32, %x : i32, %y : i32) {
+  // CHECK:    llvm.call spir_funccc @__assert_fail
+  %0 = triton_gen.2Dblockload %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=8, tile_width=8, tile_height=8, v_blocks=1, transpose=false, vnni_transform=false, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<2xi16>
+  llvm.return
+}
+}
+
+// -----
+
+llvm.func @triton_gen.2Dblockprefetch(%ptr : !llvm.ptr<1>, %base_width : i32, %base_height : i32, %base_pitch : i32, %x : i32, %y : i32) {
+  // CHECK:    llvm.call spir_funccc @__assert_fail
+  triton_gen.2Dblockprefetch %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=8, tile_width=32, tile_height=16, v_blocks=1, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32)
+  llvm.return
+}
+
+// -----
+
+module attributes {"ttg.threads-per-warp" = 16 : i32} {
+llvm.func @triton_gen.2Dblockstore(%ptr : !llvm.ptr<1>, %base_width : i32, %base_height : i32, %base_pitch : i32, %x : i32, %y : i32, %stored_val : vector<8xi16>) {
+  // CHECK:    llvm.call spir_funccc @__assert_fail
+  triton_gen.2Dblockstore %ptr, %base_width, %base_height, %base_pitch, %x, %y, %stored_val {elem_size_in_bits=8, tile_width=32, tile_height=8, v_blocks=1, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi16>)
+  llvm.return
+}
+}

third_party/intel/include/TritonGENToLLVM/TritonGENToLLVMPass.h

@@ -21,8 +21,13 @@ namespace triton {
 #define GEN_PASS_DECL
 #include "intel/include/TritonGENToLLVM/Passes.h.inc"
 
-void populateTritonGENToLLVMConversionPatterns(LLVMTypeConverter &converter,
-                                               RewritePatternSet &patterns);
+namespace gpu::intel {
+class LibCallEmitter;
+} // namespace gpu::intel
+
+void populateTritonGENToLLVMConversionPatterns(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns,
+    const mlir::triton::gpu::intel::LibCallEmitter &emitter);
 
 void registerConvertTritonGENToLLVMInterface(DialectRegistry &registry);