[TEST] Enhance test_block_io.py (#4907)

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

python/test/unit/intel/test_block_store.py renamed to python/test/unit/intel/test_block_io.py

@@ -118,9 +118,10 @@ def warps_per_cta(layout):
 @pytest.mark.parametrize("M, N", [[M, N] for M, N in itertools.product([32, 64, 128], [64, 128])])
 @pytest.mark.parametrize("dtype_str", ["float32", "float16", "int8"])
 @pytest.mark.parametrize("layout", layouts)
-@pytest.mark.parametrize("block_ptr", [True, False])
+@pytest.mark.parametrize("load_block_ptr, store_block_ptr", [(True, True), (False, False), (True, False),
+                                                             (False, True)])
 @pytest.mark.skipif(not is_xpu(), reason="Block store tests are specific to the XPU backend")
-def test_block_store(M, N, dtype_str, layout, block_ptr, device, tmp_path: pathlib.Path):
+def test_block_io(M, N, dtype_str, layout, load_block_ptr, store_block_ptr, device, tmp_path: pathlib.Path):
 
     warps = warps_per_cta(layout)
     num_warps = int(np.prod(warps))
@@ -131,41 +132,50 @@ def test_block_store(M, N, dtype_str, layout, block_ptr, device, tmp_path: pathl
 
     support_block_io = torch.xpu.get_device_capability()['has_subgroup_2d_block_io']
 
-    if block_ptr:
+    if load_block_ptr:
+        load_ops = f"""
+            %src_ptr = tt.make_tensor_ptr %src, [%M_i64, %N_i64], [%N_i64, %c1_i64], [%c0_i32, %c0_i32] {{order = array<i32: 1, 0>}} : <tensor<{M}x{N}x{ty}, #layout>>
+            %store_val = tt.load %src_ptr {{ttig.block_io = "row_major", boundaryCheck = array<i32: 0, 1>, padding = 1 : i32}} : !tt.ptr<tensor<{M}x{N}x{ty}, #layout>>
+            """
+    else:
+        load_ops = f"""
+            %src_base = tt.splat %src : !tt.ptr<{ty}> -> tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
+            %src_ptr = tt.addptr %src_base, %row_major_off : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>, tensor<{M}x{N}xi32, #layout>
+            %store_val = tt.load %src_ptr {{ttig.block_io = "row_major"}} : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
+            """
+    if store_block_ptr:
         store_ops = f"""
-            %M_i64 = arith.constant {M} : i64
-            %N_i64 = arith.constant {N} : i64
-            %c1_i64 = arith.constant 1 : i64
-            %c0_i32 = arith.constant 0 : i32
-
             %blk_ptr = tt.make_tensor_ptr %dst, [%M_i64, %N_i64], [%N_i64, %c1_i64], [%c0_i32, %c0_i32] {{order = array<i32: 1, 0>}} : <tensor<{M}x{N}x{ty}, #layout>>
             tt.store %blk_ptr, %store_val {{ttig.block_io = "row_major", boundaryCheck = array<i32: 0, 1>}} : !tt.ptr<tensor<{M}x{N}x{ty}, #layout>>
             """
     else:
         store_ops = f"""
-            %12 = tt.splat %dst : !tt.ptr<{ty}> -> tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
-            %13 = tt.addptr %12, %8 : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>, tensor<{M}x{N}xi32, #layout>
-            tt.store %13, %store_val {{ttig.block_io = "row_major"}} : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
+            %dst_base = tt.splat %dst : !tt.ptr<{ty}> -> tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
+            %dst_ptr = tt.addptr %dst_base, %row_major_off : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>, tensor<{M}x{N}xi32, #layout>
+            tt.store %dst_ptr, %store_val {{ttig.block_io = "row_major"}} : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
             """
 
     ir = f"""
     #layout = {layout}
     module attributes {{{"ttig.support_sg_2d_block," if support_block_io else ""} "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = {num_warps} : i32, ttg.target = "xpu", "ttg.threads-per-warp" = {threads_per_warp} : i32}} {{
         tt.func public @block_store(%src: !tt.ptr<{ty}> {{tt.divisibility = 16 : i32}}, %dst: !tt.ptr<{ty}> {{tt.divisibility = 16 : i32}}) {{
 
-            %stride = arith.constant dense<{N}> : tensor<{M}x1xi32, #layout>
+            %M_i64 = arith.constant {M} : i64
+            %N_i64 = arith.constant {N} : i64
+            %c1_i64 = arith.constant 1 : i64
+            %c0_i32 = arith.constant 0 : i32
+
+            %stride_N = arith.constant dense<{N}> : tensor<{M}x1xi32, #layout>
             %1 = tt.make_range {{end = {M} : i32, start = 0 : i32}} : tensor<{M}xi32, #ttg.slice<{{dim = 1, parent = #layout}}>>
             %2 = tt.expand_dims %1 {{axis = 1 : i32}} : tensor<{M}xi32, #ttg.slice<{{dim = 1, parent = #layout}}>> -> tensor<{M}x1xi32, #layout>
-            %3 = arith.muli %2, %stride : tensor<{M}x1xi32, #layout>
+            %row_stride = arith.muli %2, %stride_N : tensor<{M}x1xi32, #layout>
             %4 = tt.make_range {{end = {N} : i32, start = 0 : i32}} : tensor<{N}xi32, #ttg.slice<{{dim = 0, parent = #layout}}>>
             %5 = tt.expand_dims %4 {{axis = 0 : i32}} : tensor<{N}xi32, #ttg.slice<{{dim = 0, parent = #layout}}>> -> tensor<1x{N}xi32, #layout>
-            %6 = tt.broadcast %3 : tensor<{M}x1xi32, #layout> -> tensor<{M}x{N}xi32, #layout>
+            %6 = tt.broadcast %row_stride : tensor<{M}x1xi32, #layout> -> tensor<{M}x{N}xi32, #layout>
             %7 = tt.broadcast %5 : tensor<1x{N}xi32, #layout> -> tensor<{M}x{N}xi32, #layout>
-            %8 = arith.addi %6, %7 : tensor<{M}x{N}xi32, #layout>
-            %9 = tt.splat %src : !tt.ptr<{ty}> -> tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
-            %10 = tt.addptr %9, %8 : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>, tensor<{M}x{N}xi32, #layout>
-            %store_val = tt.load %10 : tensor<{M}x{N}x!tt.ptr<{ty}>, #layout>
+            %row_major_off = arith.addi %6, %7 : tensor<{M}x{N}xi32, #layout>
 
+            {load_ops}
             {store_ops}
 
             tt.return
@@ -181,7 +191,7 @@ def test_block_store(M, N, dtype_str, layout, block_ptr, device, tmp_path: pathl
 
     x = torch.empty_like(a)
 
-    temp_file = tmp_path / "test_block_store.ttgir"
+    temp_file = tmp_path / "test_block_io.ttgir"
     temp_file.write_text(ir)
     kernel = triton.compile(str(temp_file))
 

scripts/skiplist/lts/intel.txt

@@ -1,2 +1,2 @@
 python/test/unit/intel/test_block_load.py::test_block_load_dpas_layout
-python/test/unit/intel/test_block_store.py::test_block_store
+python/test/unit/intel/test_block_io.py::test_block_io

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -2119,6 +2119,8 @@ struct LoadOpToBlockIOConversion
       return failure();
     numOperandsPer2DLoadN =
         std::min(numOperandsPer2DLoadN, MAX_WIDTH / totalBytesPerRowPerDPASOp);
+    // vBlocks has HW limitation of 4.
+    numOperandsPer2DLoadN = std::min(numOperandsPer2DLoadN, 4u);
 
     tileHeight = instHeight * numOperandsPer2DLoadM;
     tileWidth = instWidth;