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

Author: alexbaden

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -499,7 +499,9 @@ struct LoadOpConversion
     auto tensorType = cast<RankedTensorType>(resultType);
 
     // Only lower loadOp with dpas layout encoding.
-    if (!hasDotDpasEncoding(tensorType))
+    auto encoding = tensorType.getEncoding();
+    const bool hasDpasLayout = isa<DpasEncodingAttr>(encoding);
+    if (!hasDpasLayout && !hasDotDpasEncoding(tensorType))
       return failure();
 
     Attribute blockIOAttr =
@@ -514,8 +516,8 @@ struct LoadOpConversion
            "Only row_major or column_major is supported");
     const bool memoryRowMajor = (memoryLayoutInfo == "row_major");
 
-    DotOperandEncodingAttr dotLayout = getDotEncoding(tensorType).value();
-    auto dotOrder = dotLayout.getThreadOrder();
+    auto dpasLayout = hasDpasLayout ? cast<DpasEncodingAttr>(encoding) : cast<DpasEncodingAttr>(getDotEncoding(tensorType).value().getParent());
+    auto dotOrder = dpasLayout.getThreadOrder();
     size_t rank = dotOrder.size();
     const bool valueRowMajor =
         (dotOrder[rank - 2] == 1 && dotOrder[rank - 1] == 0);
@@ -524,9 +526,17 @@ struct LoadOpConversion
            "Only row_major or column_major is allowed");
     const bool isTransposeRequired = valueRowMajor ^ memoryRowMajor;
 
-    auto dpasLayout = cast<DpasEncodingAttr>(dotLayout.getParent());
+    // auto dpasLayout = cast<DpasEncodingAttr>(dotLayout.getParent());
+    auto getOpIdx = [&]() -> unsigned {
+      if (hasDpasLayout) {
+        return 2;
+      } else {
+        auto dotLayout = getDotEncoding(tensorType).value();
+        return dotLayout.getOpIdx();
+      }
+    };
 
-    const unsigned opIdx = dotLayout.getOpIdx();
+    const unsigned opIdx = getOpIdx();
     Type eltTy = tensorType.getElementType();
     const ArrayRef<int64_t> tensorShape = tensorType.getShape();
     unsigned numElems = getTotalElemsPerThread(resultType);
@@ -543,6 +553,198 @@ struct LoadOpConversion
     SmallVector<Value> multiDimWarpId =
         delinearize(rewriter, loc, warpId, warpsPerCTA, dpasOrder);
 
+#if 1
+    if (hasDpasLayout) {
+      llvm::errs() << "rewriting tensor pointer load for dpas user!\n";
+#if 1
+    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();
+    unsigned outerDimWarpNum = std::min<unsigned>(
+        warpsPerCTA[rank - 2],
+        mlir::ceil<unsigned>(tensorShape[rank - 2], repClusterShape[rank - 2]));
+    unsigned innerDimWarpNum = std::min<unsigned>(
+        warpsPerCTA[rank - 1],
+        mlir::ceil<unsigned>(tensorShape[rank - 1], repClusterShape[rank - 1]));
+    Value outerDimWarpId =
+        urem(multiDimWarpId[rank - 2], i32_val(outerDimWarpNum));
+    Value innerDimWarpId =
+        urem(multiDimWarpId[rank - 1], i32_val(innerDimWarpNum));
+    int64_t numRepOuter = numReps[1];
+    int64_t numRepInner = numReps[2];
+
+
+    std::array<unsigned, 2> replicaStride = {
+        outerDimWarpNum * repClusterShape[rank - 2],
+        innerDimWarpNum * repClusterShape[rank - 1]};
+    std::array<unsigned, 2> warpStride = {repClusterShape[rank - 2],
+                                          repClusterShape[rank - 1]};
+
+    Value dimWarpId0 = mul(outerDimWarpId, i32_val(warpStride[0]));
+    Value dimWarpId1 = mul(innerDimWarpId, i32_val(warpStride[1]));
+    Value warpId0Offset = add(dimWarpId0, offsetBaseY);
+    Value warpId1Offset = add(dimWarpId1, offsetBaseX);
+
+    llvm::errs() << "elemsPerInstr: " << elemsPerInstr[0] << ", " << elemsPerInstr[1] << "\n";
+    ArrayRef<unsigned> repCluster = dpasLayout.getRepCluster();
+    unsigned valOffset = 0;
+    
+    SmallVector<Value> unpackedLoadedVals;
+    
+    for (int m = 0; m < numRepOuter; ++m) {
+      for (int n = 0; n < numRepInner; ++n) {
+        for (int repM = 0; repM < repCluster[0]; ++repM) {
+
+          Value offsetY = add(warpId0Offset, i32_val(m * replicaStride[0] +
+                                                     repM * elemsPerInstr[0]));
+          for (int repN = 0; repN < repCluster[1]; ++repN) {
+                      llvm::errs() << "m, n, repM, repN: " << m << ", " << n << ", " << repM << ", " << repN << "\n";
+            Value offsetX =
+                add(warpId1Offset,
+                    i32_val(n * replicaStride[1] + repN * elemsPerInstr[1]));
+       
+            assert(!isTransposeRequired);
+            auto load2dOp = rewriter.create<TritonGEN::Matrix2DBlockLoadOp>(
+                  loc, load2DGenXType,
+                  /*ptr*/ base,
+                  /*base_width*/ mul(baseWidth, elemSizeInBytes),
+                  /*base_height*/ baseHeight,
+                  /*base_pitch*/ mul(pitch, elemSizeInBytes),
+                  /*x*/ trunc(i32_ty, offsetX),
+                  /*y*/ trunc(i32_ty, offsetY),
+                  /*elem_size_in_bits*/ elemSizeInBits,
+                /*tile_width*/ elemsPerInstr[1],
+                /*tile_height*/ elemsPerInstr[0],
+                  /*v_blocks*/ 1,
+                  /*transpose*/ isTransposeRequired,
+                  /*vnni_transform*/false /*
+                  (usePackedType && !isOperandA && !isTransposeRequired &&
+                  eltTy.getIntOrFloatBitWidth() != 32)*/);
+            if (failed(load2dOp.verify())) {
+            // Explicitly invoke verifier because `triton_gen` ops are
+            // immediately lowered further to a builtin call.
+            return failure();
+          }
+
+
+          #if 0
+          llvm::errs() << "elemsPerLane: " << elemsPerLane << "\n";
+          SmallVector<int32_t> indices(elemsPerLane);
+          for (int elemIdx = 0; elemIdx < elemsPerLane;
+                    ++elemIdx) {
+                indices[elemIdx] = elemIdx * n;
+              }
+
+        #if 0  
+          llvm::errs() << "indices: ";
+          for (size_t i = 0; i < indices.size(); i++) {
+            llvm::errs() << " " << i;
+          }
+          llvm::errs() << "\n";
+          #endif 
+          DenseI32ArrayAttr attr = rewriter.getDenseI32ArrayAttr(indices);
+                Value loadVal = rewriter.create<LLVM::ShuffleVectorOp>(
+                    loc, load2DGenXType, load2dOp, load2dOp, attr);
+          #endif 
+          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);
+          }
+
+          // loadVals[{outer * packedRowNum * numLoadPerOutRepCluster +
+          //                       rep * packedRowNum + row,
+          //                   k + vblk * packedColNumPerVBlock + col}] =
+          //             bitcast(loadVal, unpackedDPASOperandType);
+#if 0
+            Value storeVal = rewriter.create<LLVM::UndefOp>(
+                loc, LLVM::getFixedVectorType(typeConverter->convertType(eltTy),
+                                              elemsPerLane));
+            for (size_t i = 0; i < elemsPerLane; ++i) {
+              storeVal = insert_element(storeVal, vals[valOffset], i32_val(i));
+              ++valOffset;
+            }
+
+            auto newOp = rewriter.create<TritonGEN::Matrix2DBlockStoreOp>(
+                loc,
+                /*ptr*/ base,
+                /*base_width*/ baseWidth,
+                /*base_height*/ height,
+                /*base_pitch*/ basePitch,
+                /*x*/ trunc(i32_ty, offsetX),
+                /*y*/ trunc(i32_ty, offsetY),
+                /*elem_size_in_bits*/ elemSizeInBits,
+                /*tile_width*/ elemsPerInstr[1],
+                /*tile_height*/ elemsPerInstr[0],
+                /*v_blocks*/ 1,
+                /*stored_val*/ bitcast(storeVal, store2DGenXType));
+
+            if (failed(newOp.verify())) {
+              // Explicitly invoke verifier because `triton_gen` ops are
+              // immediately lowered further to a builtin call.
+              return failure();
+            }
+            #endif 
+          }
+        }
+      }
+    }
+#if 0
+    return failure();
+#else
+    TritonGPUToLLVMTypeConverter *typeConverter = getTypeConverter();
+    Type llvmResultStructTy = typeConverter->convertType(op.getType());
+    Value resultStruct = packLLElements(loc, typeConverter, unpackedLoadedVals,
+                                        rewriter, llvmResultStructTy);
+    rewriter.replaceOp(op, {resultStruct});
+
+    return success();
+#endif 
+#else
+
+      ValueTable loadVals;
+
+      unsigned numRepOuter = numReps[2];
+      // TODO: calculate this instead of guessing  
+      numOperandsPer2DLoadM = 1;
+      numOperandsPer2DloadN = 1;
+      for (int outer = 0; outer < numRepOuter; ++outer) {
+        for (int rep = 0; rep < numLoadPerOutRepCluster; ++rep) {
+         for (int k = 0; k < numRepInner; k += numOperandsInnerDimPerLoad) {
+            llvm::errs() << "load: " << outer << ", " << rep << ", " << k << "\n";
+         }
+        }
+      }
+      return failure();
+#endif
+    }
+#endif 
+
     bool isOperandA = (opIdx == 0);
     SmallVector<unsigned> dpasInstShape = isOperandA
                                               ? dpasLayout.getDPASInstShapeA()
@@ -586,6 +788,7 @@ struct LoadOpConversion
 
     Type packedDPASOperandType = LLVM::getFixedVectorType(
         loadResultElemType, packedElemsPerLanePerDPASInst);
+    llvm::errs() << "packed DPAS operand type: " << packedDPASOperandType << "\n";
 
     // Outer dim: Dim M or N. Inner dim: Dim K.
     // Round the warp id fit into the tensor shape.
@@ -705,9 +908,11 @@ struct LoadOpConversion
     Value elemSizeInBytes = i32_val(originalElemBits / 8);
 
     ValueTable loadVals;
+    llvm::errs() << "Generating 2D block load for op: " << opIdx << "\n";
     for (int outer = 0; outer < numRepOuter; ++outer) {
       for (int rep = 0; rep < numLoadPerOutRepCluster; ++rep) {
         for (int k = 0; k < numRepInner; k += numOperandsInnerDimPerLoad) {
+          llvm::errs() << "outer, rep, k = " << outer << ", " << rep << ", " << k << "\n";
           Value offsetX, offsetY;
           if (opIdx == 0) {
             // A