[intel] 2Dblock runtime HW checks

Author: januszjah

python/test/unit/intel/test_block_load.py

@@ -207,3 +207,47 @@ 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)
+
+def test_block_load_asserts(M, N, dtype_str, transpose, device, tmp_path: pathlib.Path):
+    TTGIR = 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
+
+            %0 = 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>
+            tt.return
+        }
+    }
+    """
+
+    tmp = pathlib.Path("test_block_load_asserts.ttgir").resolve()
+    tmp.write_text(TTGIR)
+
+    path = str(tmp)
+    kernel = triton.compile(path)
+
+    a = torch.randn((256, 64), dtype=torch.float32, device="xpu")
+
+    addr = int(a.data_ptr())
+    
+    import ctypes
+    addr64 = ctypes.c_int64(addr).value
+
+    print(kernel.asm['spvdis'])
+    
+    # TODO catch the assert from __assert_fail
+    kernel[(1,1,1)](addr64, 64, 64, 1, 0, 0)

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

@@ -0,0 +1,29 @@
+// 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.intr.trap"() : () -> ()
+  %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.intr.trap"() : () -> ()
+  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.intr.trap"() : () -> ()
+  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
+}
+}
+
+// TODO: change checks to use __assert_fail

third_party/intel/lib/TritonGENToLLVM/TritonGENToLLVMPass.cpp

@@ -503,6 +503,82 @@ createGenISA2DBlockPrefetch(TritonGEN::Matrix2DBlockPrefetchOp op,
                                          intel::noUnwindWillReturnAttrs);
 }
 
+template <typename OpTy>
+static void validateMatrix2DBlockParameters(OpTy op,
+                                            mlir::ConversionPatternRewriter &rewriter) {
+  using namespace mlir;
+  using namespace mlir::LLVM;
+
+  Location loc = op->getLoc();
+  auto b = TritonLLVMOpBuilder(loc, rewriter);
+  MLIRContext *ctx = rewriter.getContext();
+
+  Value baseWidth = op.getBaseWidth();
+  Value baseHeight = op.getBaseHeight();
+  Value basePitch = op.getBasePitch();
+  Value x = op.getX();
+  unsigned elemSize = op.getElemSizeInBits() / 8;
+
+  if (!baseWidth.getType().isInteger(32))
+    baseWidth = rewriter.create<ZExtOp>(loc, rewriter.getI32Type(), baseWidth);
+  if (!baseHeight.getType().isInteger(32))
+    baseHeight = rewriter.create<ZExtOp>(loc, rewriter.getI32Type(), baseHeight);
+  if (!basePitch.getType().isInteger(32))
+    basePitch = rewriter.create<ZExtOp>(loc, rewriter.getI32Type(), basePitch);
+  if (!x.getType().isInteger(32))
+    x = rewriter.create<ZExtOp>(loc, rewriter.getI32Type(), x);
+
+  Value c0 = b.i32_val(0);
+  Value c4 = b.i32_val(4);
+  Value c64 = b.i32_val(64);
+  Value c24m1 = b.i32_val((1u<<24) - 1);  // 2^24 - 1
+  Value cElemSize = b.i32_val(elemSize);
+
+  // ===== validation predicates =====
+  
+  // width!=0 && width<2^24 && width%4==0
+  Value wZero = rewriter.create<ICmpOp>(loc, ICmpPredicate::eq, baseWidth, c0);
+  Value wTooLarge = rewriter.create<ICmpOp>(loc, ICmpPredicate::ugt, baseWidth, c24m1);
+  Value wRem = rewriter.create<URemOp>(loc, baseWidth, c4);
+  Value wNotAligned = rewriter.create<ICmpOp>(loc, ICmpPredicate::ne, wRem, c0);
+  Value badWidth = rewriter.create<OrOp>(loc, wZero, 
+                     rewriter.create<OrOp>(loc, wTooLarge, wNotAligned));
+  
+  // height!=0 && height<2^24
+  Value hZero = rewriter.create<ICmpOp>(loc, ICmpPredicate::eq, baseHeight, c0);
+  Value hTooLarge = rewriter.create<ICmpOp>(loc, ICmpPredicate::ugt, baseHeight, c24m1);
+  Value badHeight = rewriter.create<OrOp>(loc, hZero, hTooLarge);
+  
+  // pitch >= 64
+  Value badPitch = rewriter.create<ICmpOp>(loc, ICmpPredicate::ult, basePitch, c64);
+  
+  // x*elemSize % 4 == 0
+  Value offsetBytes = rewriter.create<MulOp>(loc, x, cElemSize);
+  Value offsetRem = rewriter.create<URemOp>(loc, offsetBytes, c4);
+  Value badOffset = rewriter.create<ICmpOp>(loc, ICmpPredicate::ne, offsetRem, c0);
+  
+  // assert on any
+  Value anyBad = rewriter.create<OrOp>(loc, badWidth,
+                  rewriter.create<OrOp>(loc, badHeight,
+                    rewriter.create<OrOp>(loc, badPitch, badOffset)));
+
+  Block *curBlock = rewriter.getBlock();
+  auto ip = rewriter.getInsertionPoint();
+  Block *contBlock = rewriter.splitBlock(curBlock, ip);
+  Region *region = contBlock->getParent();
+  Block *trapBlock = rewriter.createBlock(region, Region::iterator(contBlock));
+
+  // TODO: use __assert_fail instead of llvm.intr.trap
+  rewriter.setInsertionPointToStart(trapBlock);
+  rewriter.create<Trap>(loc);
+  rewriter.create<UnreachableOp>(loc);
+
+  rewriter.setInsertionPointToEnd(curBlock);
+  rewriter.create<CondBrOp>(loc, anyBad, trapBlock, ValueRange{}, contBlock, ValueRange{});
+
+  rewriter.setInsertionPointToStart(contBlock);
+}
+
 namespace {
 
 //===----------------------------------------------------------------------===//
@@ -636,6 +712,8 @@ struct TritonMatrix2DBlockLoadLowering
   LogicalResult
   matchAndRewrite(TritonGEN::Matrix2DBlockLoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    validateMatrix2DBlockParameters(op, rewriter);
+
     if (!isSPVBuiltinAvailable(op)) {
       // Fallback to GenISA interface.
       rewriter.replaceOp(op, createGenISA2DBlockRead(op, rewriter));
@@ -711,6 +789,8 @@ struct TritonMatrix2DBlockStoreLowering
   LogicalResult
   matchAndRewrite(TritonGEN::Matrix2DBlockStoreOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    validateMatrix2DBlockParameters(op, rewriter);
+
     if (!isSPVBuiltinAvailable(op)) {
       // Fallback to GenISA interface.
       rewriter.replaceOp(op, createGenISA2DBlockWrite(op, rewriter));
@@ -785,6 +865,8 @@ struct TritonMatrix2DBlockPrefetchLowering
   LogicalResult
   matchAndRewrite(TritonGEN::Matrix2DBlockPrefetchOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    validateMatrix2DBlockParameters(op, rewriter);
+
     if (!isSPVBuiltinAvailable(op)) {
       // Fallback to GenISA interface.
       rewriter.replaceOp(op, createGenISA2DBlockPrefetch(op, rewriter));