[TritonGEN] Lower to GenISA for 2d_block_read_transpose_32b_8r[2|4|8]x1c (#5083)

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

Author: whitneywhtsang

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

@@ -638,6 +638,54 @@ llvm.func @triton_gen.2Dblockload(%ptr : !llvm.ptr<1>, %base_width : i32, %base_
 
 // -----
 
+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:    [[ELEM_BITS:%.*]] = llvm.mlir.constant(32 : 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(true) : i1
+  // CHECK:    [[VNNI:%.*]] = llvm.mlir.constant(false) : i1
+  // CHECK:    llvm.call spir_funccc @llvm.genx.GenISA.LSC2DBlockRead.v1i32({{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, [[ELEM_BITS]], [[TILE_WIDTH]], [[TILE_HEIGHT]], [[VBLOCKS]], [[TRANSPOSE]], [[VNNI]], {{.*}})
+  %0 = triton_gen.2Dblockload %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=32, tile_width=2, tile_height=8, v_blocks=1, transpose=true, vnni_transform=false, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<1xi32>
+  llvm.return
+}
+}
+
+// -----
+
+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:    [[ELEM_BITS:%.*]] = llvm.mlir.constant(32 : i32) : i32
+  // CHECK:    [[TILE_WIDTH:%.*]] = llvm.mlir.constant(4 : i32) : i32
+  // CHECK:    [[TILE_HEIGHT:%.*]] = llvm.mlir.constant(8 : i32) : i32
+  // CHECK:    [[VBLOCKS:%.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK:    [[TRANSPOSE:%.*]] = llvm.mlir.constant(true) : i1
+  // CHECK:    [[VNNI:%.*]] = llvm.mlir.constant(false) : i1
+  // CHECK:    llvm.call spir_funccc @llvm.genx.GenISA.LSC2DBlockRead.v2i32({{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, [[ELEM_BITS]], [[TILE_WIDTH]], [[TILE_HEIGHT]], [[VBLOCKS]], [[TRANSPOSE]], [[VNNI]], {{.*}})
+  %0 = triton_gen.2Dblockload %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=32, tile_width=4, tile_height=8, v_blocks=1, transpose=true, vnni_transform=false, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<2xi32>
+  llvm.return
+}
+}
+
+// -----
+
+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:    [[ELEM_BITS:%.*]] = llvm.mlir.constant(32 : i32) : i32
+  // CHECK:    [[TILE_WIDTH:%.*]] = llvm.mlir.constant(8 : i32) : i32
+  // CHECK:    [[TILE_HEIGHT:%.*]] = llvm.mlir.constant(8 : i32) : i32
+  // CHECK:    [[VBLOCKS:%.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK:    [[TRANSPOSE:%.*]] = llvm.mlir.constant(true) : i1
+  // CHECK:    [[VNNI:%.*]] = llvm.mlir.constant(false) : i1
+  // CHECK:    llvm.call spir_funccc @llvm.genx.GenISA.LSC2DBlockRead.v4i32({{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, [[ELEM_BITS]], [[TILE_WIDTH]], [[TILE_HEIGHT]], [[VBLOCKS]], [[TRANSPOSE]], [[VNNI]], {{.*}})
+  %0 = triton_gen.2Dblockload %ptr, %base_width, %base_height, %base_pitch, %x, %y {elem_size_in_bits=32, tile_width=8, tile_height=8, v_blocks=1, transpose=true, vnni_transform=false, cache_control=Default} : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<4xi32>
+  llvm.return
+}
+}
+
+// -----
+
 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.mlir.constant(4 : i32) : i32

third_party/intel/lib/TritonGENToLLVM/TritonGENToLLVMPass.cpp

@@ -180,6 +180,21 @@ static bool isSPVBuiltinAvailable(TritonGEN::Matrix2DBlockLoadOp op) {
       op.getTileWidth() == 8 && op.getVBlocks() == 4 && !op.getVnniTransform())
     return false;
 
+  // intel_sub_group_2d_block_read_transpose_32b_8r2x1c
+  if (op.getElemSizeInBits() == 32 && op.getTileHeight() == 8 &&
+      op.getTileWidth() == 2 && op.getVBlocks() == 1 && op.getTranspose())
+    return false;
+
+  // intel_sub_group_2d_block_read_transpose_32b_8r4x1c
+  if (op.getElemSizeInBits() == 32 && op.getTileHeight() == 8 &&
+      op.getTileWidth() == 4 && op.getVBlocks() == 1 && op.getTranspose())
+    return false;
+
+  // intel_sub_group_2d_block_read_transpose_32b_8r8x1c
+  if (op.getElemSizeInBits() == 32 && op.getTileHeight() == 8 &&
+      op.getTileWidth() == 8 && op.getVBlocks() == 1 && op.getTranspose())
+    return false;
+
   // FIXME: The SPV block load only support subgroup size 16.
   int subGroupSize = triton::gpu::TritonGPUDialect::getThreadsPerWarp(
       op->getParentOfType<mlir::ModuleOp>());