[TritonGEN] Lower to GenISA for 2d_block_prefetch_16b_2r8x1c (#4649)

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

Author: whitneywhtsang

test/TritonGEN/tritongen-2Dblockprefetch-to-llvm.mlir

@@ -59,6 +59,34 @@ llvm.func @triton_gen.2Dblockprefetch(%ptr : !llvm.ptr<1>, %base_width : i32, %b
 
 // -----
 
+llvm.func @triton_gen.2Dblockprefetch(%ptr : !llvm.ptr<1>, %base_width : i32, %base_height : i32, %base_pitch : i32, %x : i32, %y : i32) {
+  // CHECK:    [[ONE0:%.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK:    [[PTR:%.*]] = llvm.ptrtoint %arg0 : !llvm.ptr<1> to i64
+  // CHECK:    [[VAL_63:%.*]] = llvm.mlir.constant(-64 : i64) : i64
+  // CHECK:    [[VAL_64:%.*]] = llvm.and [[PTR]], [[VAL_63]] : i64
+  // CHECK:    [[VAL_65:%.*]] = llvm.inttoptr [[VAL_64]] : i64 to !llvm.ptr<1>
+  // CHECK:    [[CL:%.*]] = llvm.mlir.constant(63 : i64) : i64
+  // CHECK:    [[AND:%.*]] = llvm.and [[PTR]], [[CL]] : i64
+  // CHECK:    [[TRUNC:%.*]] = llvm.trunc [[AND]] : i64 to i32
+  // CHECK:    [[ADD:%.*]] = llvm.add %arg1, [[TRUNC]] : i32
+  // CHECK:    [[TWO:%.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK:    [[SHR:%.*]] = llvm.udiv [[TRUNC]], [[TWO]] : i32
+  // CHECK:    [[X:%.*]] = llvm.add %arg4, [[SHR]] : i32
+  // CHECK:    [[BASE_ALIGNED:%.*]] = llvm.ptrtoint [[VAL_65]] : !llvm.ptr<1> to i64
+  // CHECK:    [[BASEWIDTH:%.*]] = llvm.sub [[ADD]], [[ONE0]] : i32
+  // CHECK:    [[ELEM_BITS:%.*]] = llvm.mlir.constant(16 : i32) : i32
+  // CHECK:    [[TILE_WIDTH:%.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK:    [[TILE_HEIGHT:%.*]] = llvm.mlir.constant(8 : i32) : i32
+  // CHECK:    [[VBLOCKS:%.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK:    [[TRANSPOSE:%.*]] = llvm.mlir.constant(false) : i1
+  // CHECK:    [[VNNI:%.*]] = llvm.mlir.constant(false) : i1
+  // CHECK:    llvm.call spir_funccc @llvm.genx.GenISA.LSC2DBlockPrefetch.isVoid([[BASE_ALIGNED]], [[BASEWIDTH]], {{.*}}, [[X]], {{.*}}, [[ELEM_BITS]], [[TILE_WIDTH]], [[TILE_HEIGHT]], [[VBLOCKS]], [[TRANSPOSE]], [[VNNI]], {{.*}})
+  triton_gen.2Dblockprefetch %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=16, tile_width=2, tile_height=8, v_blocks=1, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32)
+  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.mlir.constant(2 : i32) : i32
   // CHECK:      [[ElemSize:%.*]] = llvm.mlir.constant(2 : i32) : i32

third_party/intel/lib/TritonGENToLLVM/TritonGENToLLVMPass.cpp

@@ -144,6 +144,16 @@ static bool isSPVBuiltinAvailable(TritonGEN::Matrix2DBlockLoadOp op) {
   return true;
 }
 
+static bool isSPVBuiltinAvailable(TritonGEN::Matrix2DBlockPrefetchOp op) {
+  // FIXME: The following signatures are not valid in SPV interface.
+  // intel_sub_group_2d_block_prefetch_16b_2r8x1c
+  if (op.getElemSizeInBits() == 16 && op.getTileHeight() == 8 &&
+      op.getTileWidth() == 2 && op.getVBlocks() == 1)
+    return false;
+
+  return true;
+}
+
 // HW requires base address to be 64-byte aligned. Compensate the non-64-byte
 // alignment base address by adjusting the base width and x-coordinate offset.
 template <
@@ -651,6 +661,12 @@ struct TritonMatrix2DBlockPrefetchLowering
   LogicalResult
   matchAndRewrite(TritonGEN::Matrix2DBlockPrefetchOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    if (!isSPVBuiltinAvailable(op)) {
+      // Fallback to GenISA interface.
+      rewriter.replaceOp(op, createGenISA2DBlockPrefetch(op, rewriter));
+      return success();
+    }
+
     MLIRContext *ctx = rewriter.getContext();
     Location loc = op->getLoc();
     auto b = TritonLLVMOpBuilder(loc, rewriter);