[BACKEND] Optimize the lowering of tt.load with masks (#4539)

Use the block IO hardware boundary protection capability to replace the
branch in lowering tt.load with masks.

Signed-off-by: Lu,Chengjun <[email protected]>

Author: chengjunlu

test/TritonIntelGPU/tensor-pointer-load-block-2d.mlir

@@ -187,3 +187,62 @@ module attributes {ttig.support_sg_2d_block, "ttg.num-warps" = 8 : i32} {
     tt.return
   }
 }
+
+// -----
+
+#mma = #ttig.dpas<{repeatCount = 8, systolicDepth = 8, executionSize = 16, opsPerChan = 2, threadsPerWarp = 16, warpsPerCTA = [8, 1], repCluster = [2, 2]}>
+module attributes {ttig.support_sg_2d_block, "ttg.num-warps" = 8 : i32} {
+  // CHECK-LABEL: @regular_pointer_block_io
+  tt.func public @regular_pointer_block_io(%arg0: tensor<256x64x!tt.ptr<f16>, #mma>) {
+
+    %a_mask = arith.constant dense<true> : tensor<256x64xi1, #mma>
+    %a_other = arith.constant dense<0.00e+00> : tensor<256x64xf16, #mma>
+    // CHECK-NOT: llvm.cond_br
+
+    // CHECK: %[[TOP_LEFT_MASK_BOOL_0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.struct<(i1, i1, {{.*}}
+    // CHECK: %[[TOP_LEFT_MASK_BOOL_32:.*]] = llvm.extractvalue {{.*}}[32] : !llvm.struct<(i1, i1, {{.*}}
+    // CHECK: %[[TOP_LEFT_MASK_BOOL_64:.*]] = llvm.extractvalue {{.*}}[64] : !llvm.struct<(i1, i1, {{.*}}
+    // CHECK: %[[TOP_LEFT_MASK_BOOL_96:.*]] = llvm.extractvalue {{.*}}[96] : !llvm.struct<(i1, i1, {{.*}}
+
+
+    // CHECK: %[[BLOCK_SHAPE_Y:.*]] = llvm.mlir.constant(16 : i32) : i32
+    // CHECK: %[[CST0_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[CST0_1:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[TOP_LEFT_MASK_0:.*]] = llvm.zext %[[TOP_LEFT_MASK_BOOL_0]] : i1 to i8
+    // CHECK: %[[PRED:.*]] = llvm.call spir_funccc @_Z17sub_group_shufflecj(%[[TOP_LEFT_MASK_0]], %[[CST0_1]])
+    // CHECK: %[[PRED_BOOL:.*]] =  llvm.trunc %[[PRED]] : i8 to i1
+    // CHECK: %[[BASE_Y_0:.*]] = llvm.select %[[PRED_BOOL]], %[[CST0_0]], %[[BLOCK_SHAPE_Y]] : i1, i32
+    // CHECK: %[[LOAD_0:.*]] = triton_gen.2Dblockload {{.*}}, %[[BASE_Y_0]] {elem_size_in_bits = 16, tile_width = 16, tile_height = 16, v_blocks = 2
+    // CHECK: llvm.select {{.*}}, %[[LOAD_0]], {{.*}} : i1, vector<32xf16>
+
+    // CHECK: %[[CST0_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[CST0_1:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[TOP_LEFT_MASK_1:.*]] = llvm.zext %[[TOP_LEFT_MASK_BOOL_64]] : i1 to i8
+    // CHECK: %[[PRED:.*]] = llvm.call spir_funccc @_Z17sub_group_shufflecj(%[[TOP_LEFT_MASK_1]], %[[CST0_1]])
+    // CHECK: %[[PRED_BOOL:.*]] =  llvm.trunc %[[PRED]] : i8 to i1
+    // CHECK: %[[BASE_Y_1:.*]] = llvm.select %[[PRED_BOOL]], %[[CST0_0]], %[[BLOCK_SHAPE_Y]] : i1, i32
+    // CHECK: %[[LOAD_1:.*]] = triton_gen.2Dblockload {{.*}}, %[[BASE_Y_1]]  {elem_size_in_bits = 16, tile_width = 16, tile_height = 16, v_blocks = 2
+    // CHECK: llvm.select {{.*}}, %[[LOAD_1]], {{.*}} : i1, vector<32xf16>
+
+    // CHECK: %[[CST0_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[CST0_1:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[TOP_LEFT_MASK_2:.*]] = llvm.zext %[[TOP_LEFT_MASK_BOOL_32]] : i1 to i8
+    // CHECK: %[[PRED:.*]] = llvm.call spir_funccc @_Z17sub_group_shufflecj(%[[TOP_LEFT_MASK_2]], %[[CST0_1]])
+    // CHECK: %[[PRED_BOOL:.*]] =  llvm.trunc %[[PRED]] : i8 to i1
+    // CHECK: %[[BASE_Y_2:.*]] = llvm.select %[[PRED_BOOL]], %[[CST0_0]], %[[BLOCK_SHAPE_Y]] : i1, i32
+    // CHECK: %[[LOAD_2:.*]] = triton_gen.2Dblockload {{.*}}, %[[BASE_Y_2]]  {elem_size_in_bits = 16, tile_width = 16, tile_height = 16, v_blocks = 2
+    // CHECK: llvm.select {{.*}}, %[[LOAD_2]], {{.*}} : i1, vector<32xf16>
+
+    // CHECK: %[[CST0_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[CST0_1:.*]] = llvm.mlir.constant(0 : i32) : i32
+    // CHECK: %[[TOP_LEFT_MASK_3:.*]] = llvm.zext %[[TOP_LEFT_MASK_BOOL_96]] : i1 to i8
+    // CHECK: %[[PRED:.*]] = llvm.call spir_funccc @_Z17sub_group_shufflecj(%[[TOP_LEFT_MASK_3]], %[[CST0_1]])
+    // CHECK: %[[PRED_BOOL:.*]] =  llvm.trunc %[[PRED]] : i8 to i1
+    // CHECK: %[[BASE_Y_3:.*]] = llvm.select %[[PRED_BOOL]], %[[CST0_0]], %[[BLOCK_SHAPE_Y]] : i1, i32
+    // CHECK: %[[LOAD_3:.*]] = triton_gen.2Dblockload {{.*}}, %[[BASE_Y_3]]  {elem_size_in_bits = 16, tile_width = 16, tile_height = 16, v_blocks = 2
+    // CHECK: llvm.select {{.*}}, %[[LOAD_3]], {{.*}} : i1, vector<32xf16>
+    %0 = tt.load %arg0, %a_mask, %a_other {ttig.block_io = "row_major"} : tensor<256x64x!tt.ptr<f16>, #mma>
+
+    tt.return
+  }
+}

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -1208,11 +1208,47 @@ struct LoadOpToBlockIOConversion
                  << ", loadOuter:" << loadOuter << " offset: [" << offsetM
                  << ", " << offsetN << "]");
 
-            Value pred =
-                masks.size() ? masks[{offsetM, offsetN}] : b.int_val(1, 1);
-            pred = targetInfo.shuffleIdx(rewriter, loc, pred, 0);
-            Value other_ = b.undef(load2DGenXType);
+            Value offsetY = b.i32_val(0);
+            Value pred;
+            if (llMask) {
+              assert(masks.size() && "Invalid size of the masks.");
+              pred = targetInfo.shuffleIdx(rewriter, loc,
+                                           masks[{offsetM, offsetN}], 0);
+              // We leverage the GPU block I/O hardware out-of-bound protection
+              // feature by setting the offset to an invalid value when 'pred'
+              // is false (the HW will not read out-of-bounds values). Later on,
+              // after issuing the 2d block read operation, we will select the
+              // result of the load only if the mask evaluate to true, otherwise
+              // we will use 'other'.
+              offsetY = b.select(pred, offsetY, baseHeight);
+            }
+
+            // Use the top-left address of the block to load the data.
+            Value addrElem =
+                b.bitcast(ptrs[{offsetM, offsetN}], ptr_ty(ctx, 1 /*global*/));
+            addrElem = targetInfo.shuffleIdx(rewriter, loc, addrElem, 0);
+
+            Value ret = rewriter.create<TritonGEN::Matrix2DBlockLoadOp>(
+                loc, load2DGenXType,
+                /*ptr*/ addrElem,
+                /*base_width*/ baseWidth,
+                /*base_height*/ baseHeight,
+                /*base_pitch*/ pitch,
+                /*x*/ b.i32_val(0),
+                /*y*/ offsetY,
+                /*elem_size_in_bits*/ elemSizeInBits,
+                /*tile_width*/ tileWidth,
+                /*tile_height*/ tileHeight,
+                /*v_blocks*/ vBlocks,
+                /*transpose*/ false,
+                /*vnni_transform*/
+                (usePackedType && opIdx == DpasEncodingAttr::OpIdx::OperandB &&
+                 !isTransposeRequired && originalElemBits != 32));
+
             if (others.size()) {
+              assert(masks.size() == others.size() &&
+                     "The mask value has to be provided when "
+                     "the other value is provided.");
               VectorType vecTy =
                   vec_ty(eltTy, numValuesPerLoad * packedElemsNum);
 
@@ -1241,49 +1277,8 @@ struct LoadOpToBlockIOConversion
                     }
                   }
                 }
-
-              other_ = b.bitcast(v, load2DGenXType);
-
-            } else {
-              other_ = rewriter.create<LLVM::ConstantOp>(
-                  loc, load2DGenXType, rewriter.getZeroAttr(load2DGenXType));
-            }
-
-            auto createLoadInstruction = [&]() -> SmallVector<Value, 1> {
-              // Use the top-left address of the block to load the data.
-              Value addrElem = b.bitcast(ptrs[{offsetM, offsetN}],
-                                         ptr_ty(ctx, 1 /*global*/));
-              addrElem = targetInfo.shuffleIdx(rewriter, loc, addrElem, 0);
-
-              auto load2dOp = rewriter.create<TritonGEN::Matrix2DBlockLoadOp>(
-                  loc, load2DGenXType,
-                  /*ptr*/ addrElem,
-                  /*base_width*/ baseWidth,
-                  /*base_height*/ baseHeight,
-                  /*base_pitch*/ pitch,
-                  /*x*/ b.i32_val(0),
-                  /*y*/ b.i32_val(0),
-                  /*elem_size_in_bits*/ elemSizeInBits,
-                  /*tile_width*/ tileWidth,
-                  /*tile_height*/ tileHeight,
-                  /*v_blocks*/ vBlocks,
-                  /*transpose*/ false,
-                  /*vnni_transform*/
-                  (usePackedType &&
-                   opIdx == DpasEncodingAttr::OpIdx::OperandB &&
-                   !isTransposeRequired && originalElemBits != 32));
-              return {load2dOp};
-            };
-
-            Value ret;
-            // Create a predicated load operation.
-            if (llMask) {
-              Block &endBlock = LLVM::intel::createPredicatedBlock(
-                  rewriter, loc, pred, SmallVector<Value, 1>{other_},
-                  createLoadInstruction);
-              ret = *endBlock.args_begin();
-            } else {
-              ret = createLoadInstruction()[0];
+              Value others = b.bitcast(v, load2DGenXType);
+              ret = b.select(pred, ret, others);
             }
 
             unsigned numOperandsM = opIdx != DpasEncodingAttr::OpIdx::OperandB