按地址连续性为 GPU 访存选择线程布局

84    llvm::MapVector<Operation *, Attribute> layoutMap;85    int threadsPerWarp = TritonGPUDialect::getThreadsPerWarp(moduleOp);86    moduleOp.walk([&](Operation *curr) {87      Value ptr = getMemAccessPtr(curr);88      if (!ptr)89        return;90      // We only convert `tensor<tt.ptr<>>` load/store91      bool isPtrTensor = false;92      if (auto tensorType = dyn_cast<RankedTensorType>(ptr.getType()))93        isPtrTensor = isa<PointerType>(tensorType.getElementType());94      if (!isPtrTensor)95        return;96      int numWarps = lookupNumWarps(curr);9798      auto tensorType = cast<RankedTensorType>(ptr.getType());99      CGAEncodingAttr cgaLayout = getCGALayout(tensorType.getEncoding());100      SmallVector<int64_t> shapePerCTA = getShapePerCTA(tensorType);101      auto layout =102          buildCoalescedEncoding(axisInfoAnalysis, curr, numWarps,103                                 threadsPerWarp, cgaLayout, shapePerCTA);104      layoutMap[curr] = layout;105    });

31  auto contiguity = axisInfoAnalysis.getAxisInfo(ptr)->getContiguity();32  SmallVector<unsigned> order = getOrderFromContiguity(contiguity);33  LDBG("order=[" << triton::join(order, ", ") << "]");3435  auto matchesShape = [&refTensorType](const Value &val) {36    auto rttType = dyn_cast<RankedTensorType>(val.getType());37    return rttType && rttType.getShape() == refTensorType.getShape();38  };3940  // The desired divisibility is the maximum divisibility among all dependent41  // pointers which have the same shape and order as `ptr`.42  llvm::SmallSetVector<Operation *, 32> memAccessesSameOrder;43  memAccessesSameOrder.insert(op);44  if (ptr.getDefiningOp()) {45    for (Operation *use : mlir::getSlice(op)) {46      Value val = getMemAccessPtr(use);47      if (!val || !matchesShape(val) || memAccessesSameOrder.contains(use))48        continue;49      auto currOrder = getOrderFromContiguity(50          axisInfoAnalysis.getAxisInfo(val)->getContiguity());51      if (order == currOrder) {

177unsigned getNumElementsPerThread(Operation *op, SmallVector<unsigned> order,178                                 ModuleAxisInfoAnalysis &axisInfoAnalysis,179                                 ArrayRef<int64_t> shapePerCTA) {180  Value val = getMemAccessPtr(op);181  auto ty = cast<RankedTensorType>(val.getType());182  AxisInfo &valInfo = *axisInfoAnalysis.getAxisInfo(val);183  unsigned elemNumBits = getElementBitWidth(ty);184  unsigned elemNumBytes = std::max(elemNumBits / 8, 1u);185  unsigned maxMultipleBytes = valInfo.getDivisibility(order[0]);186  unsigned maxMultiple = std::max(maxMultipleBytes / elemNumBytes, 1u);187  unsigned maxContig =188      std::min(valInfo.getContiguity(order[0]), shapePerCTA[order[0]]);189  unsigned alignment = std::min(maxMultiple, maxContig);190  unsigned maxElementsPerThread = getMaxElementsPerThread(op);191  unsigned currPerThread = std::min(alignment, maxElementsPerThread);192  LDBG("elemNumBytes: " << elemNumBytes193                        << ", divisibility: " << maxMultipleBytes194                        << ", contig: " << valInfo.getContiguity(order[0])195                        << ", maximum: " << maxElementsPerThread196                        << ", alignment: " << alignment);197  return currPerThread;198}

76  perThread = std::min<int>(perThread, std::max(numElems / numThreads, 1));77  LDBG("perThread: " << perThread);7879  if (!dyn_cast<triton::LoadOp>(op)) {80    // For ops that can result in a global memory write, we should enforce81    // that each thread handles at most 128 bits, which is the widest82    // available vectorized store op; otherwise, the store will have "gaps"83    // in the memory write at the warp level, resulting in worse performance.84    // For loads, we can expect that the gaps won't matter due to the L185    // cache.86    perThread = std::min<int>(87        perThread,88        getNumElementsPerThread(op, order, axisInfoAnalysis, shapePerCTA));89  }90  SmallVector<unsigned> sizePerThread(refTensorType.getRank(), 1);91  sizePerThread[order[0]] = perThread;92  return BlockedEncodingAttr::get(op->getContext(), refTensorType.getShape(),93                                  sizePerThread, order, numWarps,94                                  threadsPerWarp, cgaLayout);

1336Operation *convertDistributedOpEncoding(Attribute encoding, Operation *op) {1337  OpBuilder builder(op);1338  // Convert operands1339  SmallVector<Value, 4> newArgs;1340  for (auto &opOperand : op->getOpOperands()) {1341    Value operand = opOperand.get();1342    if (isDistributedOpEncodingOperand(opOperand)) {1343      Type newType = getNewType(operand.getType(), encoding);1344      newArgs.push_back(triton::gpu::ConvertLayoutOp::create(1345          builder, op->getLoc(), newType, operand));1346    } else {1347      newArgs.push_back(operand);1348    }1349  }13501351  // Convert output types1352  SmallVector<Type, 4> newTypes;1353  for (auto t : op->getResultTypes()) {1354    bool isAsync = isa<triton::gpu::AsyncCopyGlobalToLocalOp>(op);1355    newTypes.push_back(isAsync ? t : getNewType(t, encoding));1356  }

1358  // Construct new op with the new encoding1359  Operation *newOp = builder.create(op->getLoc(), op->getName().getIdentifier(),1360                                    newArgs, newTypes, op->getAttrs());13611362  // Cast the results back to the original layout1363  for (size_t i = 0; i < op->getNumResults(); i++) {1364    Value newResult = newOp->getResult(i);1365    if (newTypes[i] != op->getResultTypes()[i]) {1366      newResult = triton::gpu::ConvertLayoutOp::create(1367          builder, op->getLoc(), op->getResult(i).getType(), newResult);1368    }1369    op->getResult(i).replaceAllUsesWith(newResult);1370  }1371  op->erase();1372  return newOp;