Support 2D block IO to load column major for DotOp matrix

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

Author: chengjunlu

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -1648,6 +1648,15 @@ struct LoadOpConversion
       usePackedType = true;
     }
 
+    if (isTransposeRequired) {
+      if (!usePackedType) {
+        // use the d32 transpose 2d load.
+        loadResultElemType = i32_ty;
+        packedElemsPerLanePerDPASInst = 32 / elemSizeInBits;
+        usePackedType = true;
+      }
+    }
+
     Type packedDPASOperandType =
         LLVM::getVectorType(loadResultElemType, packedElemsPerLanePerDPASInst);
 
@@ -2082,12 +2091,14 @@ struct LoadOpConversion
             offsetX = b.udiv(offsetX, b.i32_val(32 / originalElemBits));
           }
 
+          Value base_width = b.mul(baseWidth, elemSizeInBytes);
+          Value base_pitch = b.mul(pitch, elemSizeInBytes);
           auto load2dOp = rewriter.create<TritonGEN::Matrix2DBlockLoadOp>(
               loc, load2DGenXType,
               /*ptr*/ base,
-              /*base_width*/ b.mul(baseWidth, elemSizeInBytes),
+              /*base_width*/ base_width,
               /*base_height*/ baseHeight,
-              /*base_pitch*/ b.mul(pitch, elemSizeInBytes),
+              /*base_pitch*/ base_pitch,
               /*x*/ b.trunc(i32_ty, offsetX),
               /*y*/ b.trunc(i32_ty, offsetY),
               /*elem_size_in_bits*/ elemSizeInBits,
@@ -2105,6 +2116,10 @@ struct LoadOpConversion
             rewriter.eraseOp(load2dOp);
             return failure();
           }
+#if 0
+          targetInfo.printf(rewriter, "base: %p, baseWidth: %d, baseHeight:%d, pitch:%d, offset_x:%d, offset_y:%d, loadVal: %d",
+            {base, base_width, baseHeight, base_pitch, offsetX, offsetY, load2dOp.getResult()});
+#endif
           LLVM_DEBUG(llvm::dbgs() << "Generated load op: " << load2dOp << "\n");
 
           unsigned packedRowNum = opIdx == DpasEncodingAttr::OpIdx::OperandA
@@ -2166,11 +2181,14 @@ struct LoadOpConversion
                                           vblk * packedColNumPerVBlock + col)
                         << ", " << std::to_string(k + row) << "\n";
                   });
+                  auto ret = b.bitcast(loadVal, unpackedDPASOperandType);
+#if 0
+                  targetInfo.printf(rewriter, "loadVal: %d", {ret});
+#endif
                   loadVals[{outer * packedColNum * numLoadPerOutRepCluster +
                                 rep * packedColNum +
                                 vblk * packedColNumPerVBlock + col,
-                            k + row}] =
-                      b.bitcast(loadVal, unpackedDPASOperandType);
+                            k + row}] = ret;
                 } break;
                 case DpasEncodingAttr::OpIdx::OperandC: {
                   llvm_unreachable("unexpected OpIdx::OperandC");