Use block loads for post-dpas vector computation 3/?

Author: alexbaden

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -529,7 +529,6 @@ struct LoadOpConversion
            "Only row_major or column_major is allowed");
     const bool isTransposeRequired = valueRowMajor ^ memoryRowMajor;
 
-    // auto dpasLayout = cast<DpasEncodingAttr>(dotLayout.getParent());
     auto getOpIdx = [&]() -> unsigned {
       if (hasDpasLayout) {
         return 2;
@@ -541,6 +540,8 @@ struct LoadOpConversion
 
     const unsigned opIdx = getOpIdx();
     Type eltTy = tensorType.getElementType();
+    unsigned elemSizeInBits = eltTy.getIntOrFloatBitWidth();
+
     const ArrayRef<int64_t> tensorShape = tensorType.getShape();
     unsigned numElems = getTotalElemsPerThread(resultType);
     SmallVector<int64_t> numReps =
@@ -556,30 +557,30 @@ struct LoadOpConversion
     SmallVector<Value> multiDimWarpId =
         delinearize(rewriter, loc, warpId, warpsPerCTA, dpasOrder);
 
-    // TODO: de-duplicate with above and below code
     if (hasDpasLayout) {
-      // llvm::errs() << "rewriting tensor pointer load for dpas user!\n";
+      if (isTransposeRequired) {
+        // TODO: this would likely require a shuffle to match the expected
+        // ordering coming out of the DPAS layout and requires more
+        // investigation
+        return failure();
+      }
+
       MLIRContext *ctx = rewriter.getContext();
 
-      Type eltTy = tensorType.getElementType();
-      // llvm::errs() << "Element type: " << eltTy << "\n";
-      unsigned elemSizeInBits = eltTy.getIntOrFloatBitWidth();
       Value elemSizeInBytes = i32_val(elemSizeInBits / 8);
 
       SmallVector<unsigned> elemsPerInstr = dpasLayout.getDPASInstShapeC();
       int64_t elemsPerLane = product<unsigned>(elemsPerInstr) / threadsPerWarp;
       Type load2DGenXType =
           LLVM::getFixedVectorType(IntegerType::get(ctx, elemSizeInBits),
                                    elemsPerLane); // make it opaque type.
-      // llvm::errs() << "load 2d gen x type: " << load2DGenXType << "\n";
 
       auto [base, baseWidth, baseHeight, rowStride, colStride, offsetBaseX,
             offsetBaseY] =
           getValuesFromBlockPointerStruct(adaptor.getPtr(), rewriter);
       baseWidth = trunc(i32_ty, baseWidth);
       baseHeight = trunc(i32_ty, baseHeight);
 
-      // always row order coming out of DPAS
       auto pitch = trunc(i32_ty, rowStride);
 
       SmallVector<unsigned> repClusterShape = dpasLayout.getShapeC();
@@ -626,7 +627,6 @@ struct LoadOpConversion
                   add(warpId1Offset,
                       i32_val(n * replicaStride[1] + repN * elemsPerInstr[1]));
 
-              assert(!isTransposeRequired);
               auto load2dOp = rewriter.create<TritonGEN::Matrix2DBlockLoadOp>(
                   loc, load2DGenXType,
                   /*ptr*/ base,
@@ -636,13 +636,11 @@ struct LoadOpConversion
                   /*x*/ trunc(i32_ty, offsetX),
                   /*y*/ trunc(i32_ty, offsetY),
                   /*elem_size_in_bits*/ elemSizeInBits,
-                /*tile_width*/ elemsPerInstr[1],
-                /*tile_height*/ elemsPerInstr[0],
+                  /*tile_width*/ elemsPerInstr[1],
+                  /*tile_height*/ elemsPerInstr[0],
                   /*v_blocks*/ 1,
-                  /*transpose*/ isTransposeRequired,
-                  /*vnni_transform*/false /*
-                  (usePackedType && !isOperandA && !isTransposeRequired &&
-                  eltTy.getIntOrFloatBitWidth() != 32)*/);
+                  /*transpose*/ false,
+                  /*vnni_transform*/ false);
               if (failed(load2dOp.verify())) {
                 // Explicitly invoke verifier because `triton_gen` ops are
                 // immediately lowered further to a builtin call.
@@ -651,8 +649,7 @@ struct LoadOpConversion
 
               Value ret = bitcast(
                   load2dOp, LLVM::getFixedVectorType(eltTy, elemsPerLane));
-              // llvm::errs() << "ret: " << ret << "\n";
-              // each load should give us one column
+
               for (size_t i = 0; i < elemsPerLane; i++) {
                 Value loaded = extract_element(eltTy, ret, i32_val(i));
                 unpackedLoadedVals.push_back(loaded);
@@ -701,11 +698,11 @@ struct LoadOpConversion
     // input operands to DPAS.
     // TODO: add support for int4 and int2.
     unsigned opsPerChannel = dpasLayout.getOpsPerChannel();
-    unsigned elemBits = eltTy.getIntOrFloatBitWidth();
-    if ((opsPerChannel == 4 && elemBits == 8) ||
-        (opsPerChannel == 2 && elemBits == 16) ||
-        (opsPerChannel == 1 && elemBits == 32)) {
-      loadResultElemType = (isOperandA && elemBits != 32) ? i16_ty : i32_ty;
+    if ((opsPerChannel == 4 && elemSizeInBits == 8) ||
+        (opsPerChannel == 2 && elemSizeInBits == 16) ||
+        (opsPerChannel == 1 && elemSizeInBits == 32)) {
+      loadResultElemType =
+          (isOperandA && elemSizeInBits != 32) ? i16_ty : i32_ty;
       packedElemsPerLanePerDPASInst =
           isOperandA ? elemsPerLanePerDPASInst / (opsPerChannel == 4 ? 2 : 1)
                      : elemsPerLanePerDPASInst / opsPerChannel;
@@ -779,7 +776,7 @@ struct LoadOpConversion
 
     // PVC 2D load supports 64 bytes per row at most. Load multiple dot operands
     // by enlarging the vBlocks.
-    unsigned totalBytesPerRowPerDPASOp = tileWidth * elemBits / 8;
+    unsigned totalBytesPerRowPerDPASOp = tileWidth * elemSizeInBits / 8;
     numOperandsPer2DloadN =
         std::min(numOperandsPer2DloadN, 64 / totalBytesPerRowPerDPASOp);
     vBlocks = numOperandsPer2DloadN;
@@ -823,12 +820,12 @@ struct LoadOpConversion
     baseWidth = trunc(i32_ty, baseWidth);
     baseHeight = trunc(i32_ty, baseHeight);
 
-    unsigned originalElemBits = elemBits;
+    const unsigned originalElemBits = elemSizeInBits;
     if (isTransposeRequired) {
       // adjust the block io parameter to align HW's limitations on
       // transposing load.
       tileWidth = tileWidth / (32 / originalElemBits);
-      elemBits = 32;
+      elemSizeInBits = 32;
     }
     Value elemSizeInBytes = i32_val(originalElemBits / 8);
 
@@ -872,14 +869,14 @@ struct LoadOpConversion
               /*base_pitch*/ mul(pitch, elemSizeInBytes),
               /*x*/ trunc(i32_ty, offsetX),
               /*y*/ trunc(i32_ty, offsetY),
-              /*elem_size_in_bits*/ elemBits,
+              /*elem_size_in_bits*/ elemSizeInBits,
               /*tile_width*/ tileWidth,
               /*tile_height*/ tileHeight,
               /*v_blocks*/ vBlocks,
               /*transpose*/ isTransposeRequired,
               /*vnni_transform*/
               (usePackedType && !isOperandA && !isTransposeRequired &&
-               eltTy.getIntOrFloatBitWidth() != 32));
+               originalElemBits != 32));
           if (failed(load2dOp.verify())) {
             // Explicitly invoke verifier because `triton_gen` ops are
             // immediately lowered further to a builtin call.