address review comments

Author: binarman

third_party/amd/lib/TritonAMDGPUTransforms/OptimizeBufferOpPtr.cpp

@@ -101,35 +101,20 @@ struct AdvanceBasePointer : public OpRewritePattern<scf::ForOp> {
     return Value();
   }
 
-  // load is a target buffer load
-  // offsetIncrement is a tensor added to offsets on each iteration
-  // baseIncrement is a scalar which will be added to base pointer after
-  // optimization offsetInitialized is a value of offset on first loop iteration
-  // incrementOp is an operation that advances offset tensor
-  struct LoadData {
-    Operation *load;
+  // Description of struct fields:
+  // - op is a target BufferOp, for example BufferLoadOp or BufferLoadToLocalOp
+  // - offsetIncrement is a tensor added to offsets on each iteration
+  // - baseIncrement is a scalar which will be added to base pointer
+  // - offsetInitializer is a value of offset on first loop iteration
+  // - incrementOp is an operation that advances offset tensor
+  struct BufferOpInfo {
+    amdttg::BufferOpAddressinInterface op;
     Value offsetIncrement;
     Value baseIncrement;
     Value offsetInitializer;
     Operation *incrementOp;
   };
 
-  static Value getOffset(Operation *load) {
-    if (auto specific = dyn_cast<amdttg::BufferLoadOp>(load))
-      return specific.getOffsets();
-    if (auto specific = dyn_cast<amdttg::BufferLoadToLocalOp>(load))
-      return specific.getOffsets();
-    assert(false && "unsupported operation type");
-  }
-
-  static Value getBasePtr(Operation *load) {
-    if (auto specific = dyn_cast<amdttg::BufferLoadOp>(load))
-      return specific.getPtr();
-    if (auto specific = dyn_cast<amdttg::BufferLoadToLocalOp>(load))
-      return specific.getPtr();
-    assert(false && "unsupported operation type");
-  }
-
   static void setOffset(Operation *load, Value newOffset) {
     assert((isa<amdttg::BufferLoadOp, amdttg::BufferLoadToLocalOp>(load)));
     const int offsetIdx = isa<amdttg::BufferLoadOp>(load) ? 1 : 2;
@@ -143,11 +128,11 @@ struct AdvanceBasePointer : public OpRewritePattern<scf::ForOp> {
   }
 
   // Perform series of checks to decide if given operation could be optimized.
-  // If optimization is possible, return filled LoadData
-  static std::optional<LoadData> analyzeLoad(Operation *loadOp,
-                                             scf::ForOp targetFor) {
-    LDBG("Analyzing: " << *loadOp);
-    Value maybeOffsetsBlockArg = getOffset(loadOp);
+  // If optimization is possible, return filled BufferOpInfo
+  static std::optional<BufferOpInfo>
+  analyzeBufferOp(amdttg::BufferOpAddressinInterface op, scf::ForOp targetFor) {
+    LDBG("Analyzing: " << *op);
+    Value maybeOffsetsBlockArg = op.getOffsets();
     auto maybeOffsetDefOp = maybeOffsetsBlockArg.getDefiningOp();
     if (maybeOffsetDefOp && isa<arith::AddIOp>(maybeOffsetDefOp)) {
       for (auto &use : maybeOffsetDefOp->getUses()) {
@@ -174,7 +159,7 @@ struct AdvanceBasePointer : public OpRewritePattern<scf::ForOp> {
       LDBG("Rejected: expect load offset to be a target loop argument");
       return {};
     }
-    auto basePtr = getBasePtr(loadOp);
+    auto basePtr = op.getPtr();
     auto defOpBlock = basePtr.getParentBlock();
     if (!defOpBlock->getParentOp()->isProperAncestor(targetFor)) {
       LDBG("Rejected: expect load base Ptr to be invariant to the loop");
@@ -207,33 +192,33 @@ struct AdvanceBasePointer : public OpRewritePattern<scf::ForOp> {
       return {};
     }
     Value offsetInitializer = forOp.getInitArgs()[offsetOperandNo];
-    LoadData data = {loadOp, advanceStep, Value(), offsetInitializer,
-                     incrementOp};
+    BufferOpInfo data = {op, advanceStep, Value(), offsetInitializer,
+                         incrementOp};
     LDBG("Load is suitable for offset pointer optimization");
     return data;
   }
 
   // Create scalar values which will increment load base ptr
-  // Fills appropriate fields in given LoadData structures
+  // Fills appropriate fields in given BufferOpInfo structures
   static void createScalarIncrements(PatternRewriter &rewriter,
-                                     SmallVector<LoadData> &loads) {
-    for (auto &loadData : loads) {
-      auto scalarStep = scalarizeValue(rewriter, loadData.offsetIncrement);
-      loadData.baseIncrement = scalarStep;
+                                     SmallVector<BufferOpInfo> &loads) {
+    for (auto &BufferOpInfo : loads) {
+      auto scalarStep = scalarizeValue(rewriter, BufferOpInfo.offsetIncrement);
+      BufferOpInfo.baseIncrement = scalarStep;
     }
   }
 
-  static bool isAddFirst(LoadData &ld) {
-    return getOffset(ld.load).getDefiningOp() == ld.incrementOp;
+  static bool isAddFirst(BufferOpInfo &info) {
+    return info.op.getOffsets().getDefiningOp() == info.incrementOp;
   }
 
   static scf::ForOp
   cloneLoopWithBasePtrIncrements(PatternRewriter &rewriter, scf::ForOp forOp,
-                                 SmallVector<LoadData> &loads) {
+                                 SmallVector<BufferOpInfo> &infoList) {
     // Create new loop with additional arguments
     llvm::SmallVector<Value> newLoopArgs(forOp.getInitArgs());
-    for (auto loadData : loads) {
-      newLoopArgs.push_back(getBasePtr(loadData.load));
+    for (auto info : infoList) {
+      newLoopArgs.push_back(info.op.getPtr());
     }
     rewriter.setInsertionPoint(forOp);
     auto newForOp = rewriter.create<scf::ForOp>(
@@ -251,19 +236,19 @@ struct AdvanceBasePointer : public OpRewritePattern<scf::ForOp> {
       rewriter.clone(op, mapping);
     }
     // Create base pointer increment operations
-    auto basePtrs = newBlock->getArguments().take_back(loads.size());
+    auto basePtrs = newBlock->getArguments().take_back(infoList.size());
     llvm::SmallVector<Value> nextIterBasePtrs;
-    for (auto [loadData, basePtr] : llvm::zip(loads, basePtrs)) {
-      if (isAddFirst(loadData)) {
+    for (auto [info, basePtr] : llvm::zip(infoList, basePtrs)) {
+      if (isAddFirst(info)) {
         rewriter.setInsertionPoint(newBlock, newBlock->begin());
       } else {
         rewriter.setInsertionPoint(newBlock, newBlock->end());
       }
-      Value step = loadData.baseIncrement;
-      if (mapping.contains(loadData.baseIncrement)) {
-        step = mapping.lookup(loadData.baseIncrement);
+      Value step = info.baseIncrement;
+      if (mapping.contains(info.baseIncrement)) {
+        step = mapping.lookup(info.baseIncrement);
       }
-      auto loc = loadData.incrementOp->getLoc();
+      auto loc = info.incrementOp->getLoc();
       auto ptrType = basePtr.getType();
       auto nextIterBasePtr =
           rewriter.create<triton::AddPtrOp>(loc, ptrType, basePtr, step);
@@ -282,46 +267,46 @@ struct AdvanceBasePointer : public OpRewritePattern<scf::ForOp> {
                                   newYieldOperands);
     // Replace dynamic load offsets with invariant value
     // Replace base ptr with incrementing value
-    for (auto [loadData, basePtr, nextBasePtr] :
-         llvm::zip(loads, basePtrs, nextIterBasePtrs)) {
-      auto newLoad = mapping.lookup<Operation *>(loadData.load);
-      setOffset(newLoad, loadData.offsetInitializer);
+    for (auto [info, basePtr, nextBasePtr] :
+         llvm::zip(infoList, basePtrs, nextIterBasePtrs)) {
+      auto newLoad = mapping.lookup<Operation *>(info.op.getOperation());
+      setOffset(newLoad, info.offsetInitializer);
       // two cases:
       // 1. first advance pointer, then load
       // 2. load uses pointers from loop arguments, advanced pointer used on
       // next iteration
-      Value advancingBasePtr = isAddFirst(loadData) ? nextBasePtr : basePtr;
+      Value advancingBasePtr = isAddFirst(info) ? nextBasePtr : basePtr;
       setBasePtr(newLoad, advancingBasePtr);
     }
     return newForOp;
   }
 
-  template <typename OpType>
-  static void collectLoads(SmallVector<LoadData> &loads, scf::ForOp forOp) {
-    forOp.walk([&loads, forOp](OpType loadOp) {
-      auto loadData = analyzeLoad(loadOp, forOp);
-      if (loadData.has_value()) {
-        loads.push_back(loadData.value());
+  static SmallVector<BufferOpInfo> collectBufferOps(scf::ForOp forOp) {
+    SmallVector<BufferOpInfo> list;
+    forOp.walk([&list, forOp](amdttg::BufferOpAddressinInterface op) {
+      auto info = analyzeBufferOp(op, forOp);
+      if (info.has_value()) {
+        list.push_back(info.value());
       }
     });
+    return list;
   }
 
   LogicalResult matchAndRewrite(scf::ForOp forOp,
                                 PatternRewriter &rewriter) const override {
-    LDBG("Analyzing ForOp for for offset pointer optimization: " << forOp);
-    // Gather buffer loads which could be optimized
-    SmallVector<LoadData> loads;
-    collectLoads<triton::amdgpu::BufferLoadOp>(loads, forOp);
-    collectLoads<triton::amdgpu::BufferLoadToLocalOp>(loads, forOp);
+    LDBG("Analyzing ForOp for offset pointer optimization: " << forOp);
+    // Gather buffer buffer operations which could be optimized
+    SmallVector<BufferOpInfo> infoList = collectBufferOps(forOp);
 
-    if (loads.empty())
-      return rewriter.notifyMatchFailure(forOp, "no suitable buffer loads");
+    if (infoList.empty())
+      return rewriter.notifyMatchFailure(forOp,
+                                         "no suitable buffer operations");
 
     // Perform IR transformation
-    createScalarIncrements(rewriter, loads);
-    auto newForOp = cloneLoopWithBasePtrIncrements(rewriter, forOp, loads);
-    rewriter.replaceAllUsesWith(forOp.getResults(),
-                                newForOp.getResults().drop_back(loads.size()));
+    createScalarIncrements(rewriter, infoList);
+    auto newForOp = cloneLoopWithBasePtrIncrements(rewriter, forOp, infoList);
+    rewriter.replaceAllUsesWith(
+        forOp.getResults(), newForOp.getResults().drop_back(infoList.size()));
     rewriter.eraseOp(forOp);
     return success();
   }