[ssbuffer](feat) Non-Interruptive block_id Allocation

third_party/ascend/include/DynamicCVPipeline/Common/Utils.h

@@ -121,8 +121,14 @@ bool isScfOp(Operation *op);
 bool isOnlyDirectlyUse(Operation *preOp, Operation *nextOp,
                        const CVPipeline::MemoryDependenceGraph &memGraph);
 
-inline bool isCubeOp(Operation *op) {
-  return !isScfOp(op) && CVPipeline::getOpCoreType(op) == CoreType::CUBE_ONLY;
+CoreType getCoreTypeOfSimpleOpOrCf(Operation *op);
+
+inline bool isCubeSimpleOpOrCf(Operation *op) {
+  return getCoreTypeOfSimpleOpOrCf(op) == CoreType::CUBE_ONLY;
+}
+
+inline bool isVectorSimpleOpOrCf(Operation *op) {
+  return getCoreTypeOfSimpleOpOrCf(op) == CoreType::VECTOR_ONLY;
 }
 
 bool isVectorOnlyOp(Operation *op);
@@ -150,6 +156,27 @@ bool allResultHasOneUser(Operation *op);
 
 int64_t getBTSizeFromValidBroadcastOp(linalg::BroadcastOp broadcastOp);
 
+int getLoopCarriedArgIndex(Value operand, Block *block);
+
+CoreType getValueCoreType(Value value);
+
+inline OpOperand *getTiedYieldOperand(Value value, Block *block) {
+  int argIdx = getLoopCarriedArgIndex(value, block);
+  if (argIdx == -1) {
+    return nullptr;
+  }
+  auto *terminator = block->getTerminator();
+  return &terminator->getOpOperand(argIdx);
+}
+
+inline Operation *getLoopCarriedDefOp(Value value, Block *block) {
+  auto *yieldOperand = getTiedYieldOperand(value, block);
+  if (yieldOperand && yieldOperand->get()) {
+    return yieldOperand->get().getDefiningOp();
+  }
+  return nullptr;
+}
+
 } // namespace CVPipeline
 } // namespace mlir
 

third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h

@@ -23,21 +23,58 @@
 #ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMMON_H
 #define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMMON_H
 
-#include "DynamicCVPipeline/Common/MemoryEffectsTracker.h"
-#include "DynamicCVPipeline/Common/Utils.h"
-#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Operation.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Value.h"
+
+#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h"
+
+#include "DynamicCVPipeline/Common/MemoryEffectsTracker.h"
 
 namespace mlir {
 namespace CVPipeline {
 
+class DependencyHelper {
+  using PredFn = llvm::function_ref<void(Operation *)>;
+
+  template <typename Fn>
+  static auto mapToAncestorInBlock(Block *block, Fn &&pred) {
+    return [block, pred = std::forward<Fn>(pred)](Operation *op) {
+      if (auto *ancestor = block->findAncestorOpInBlock(*op)) {
+        return pred(ancestor);
+      }
+    };
+  }
+
+public:
+  const MemoryDependenceGraph &memGraph;
+
+  explicit DependencyHelper(const MemoryDependenceGraph &memGraph)
+      : memGraph(memGraph) {}
+
+  void forEachUser(Operation *op, PredFn pred) const;
+
+  enum class SourceMode { Default, AcrossIterArg };
+  template <SourceMode SM = SourceMode::Default>
+  void forEachSource(Operation *op, PredFn pred) const;
+
+  void forEachUserInSameBlock(Operation *op, PredFn pred) const {
+    forEachUser(op, mapToAncestorInBlock(op->getBlock(), pred));
+  }
+
+  template <SourceMode ST>
+  void forEachSourceInSameBlock(Operation *op, PredFn pred) const {
+    forEachSource<ST>(op, mapToAncestorInBlock(op->getBlock(), pred));
+  }
+};
+
 Operation *getAncestorInBlock(Operation *inner, Block *block);
 void initializeIndegreeForBlock(Block *block,
                                 llvm::DenseMap<Operation *, int> &indegree,
-                                const MemoryDependenceGraph &memGraph,
+                                const DependencyHelper &depHelper,
                                 ComputeBlockIdManager &bm);
 
 } // namespace CVPipeline

third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h

@@ -23,12 +23,15 @@
 #ifndef TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMPUTE_BLOCK_ID_MANAGER_H
 #define TRITON_ADAPTER_DYNAMIC_CV_PIPELINE_PLAN_COMPUTE_BLOCK_COMPUTE_BLOCK_ID_MANAGER_H
 
-#include "mlir/IR/Operation.h"
+#include <optional>
+
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/LogicalResult.h"
-#include <mutex>
+
+#include "mlir/IR/Operation.h"
+
+#include "DynamicCVPipeline/Common/Utils.h"
 
 namespace mlir {
 namespace CVPipeline {
@@ -47,17 +50,30 @@ class ComputeBlockIdManager {
   llvm::LogicalResult markOpsWithNewId(llvm::SmallVectorImpl<Operation *> &ops);
   void updateBlockId(Operation *op, int blockId);
 
+  bool shouldInheritFromParent(Block *block, CoreType requiredCoreType) const;
+  llvm::LogicalResult inheritFromParent(Block *block);
+
   llvm::SmallVector<Operation *> getOpsByBlockId(int blockId);
-  int getBlockIdByOp(Operation *op);
-  void reset();
+
+  // Get operations that share the same block_id AND mlir block of op
+  llvm::SmallVector<Operation *> getOpsInSameBlock(Operation *op) const;
+
+  std::optional<int> getBlockIdByOpOpt(Operation *op) const;
   int getNextId();
 
+  int getBlockIdByOp(Operation *op);
+
+  ~ComputeBlockIdManager() = default;
+  ComputeBlockIdManager(const ComputeBlockIdManager &) = delete;
+  ComputeBlockIdManager &operator=(const ComputeBlockIdManager &) = delete;
+  ComputeBlockIdManager(ComputeBlockIdManager &&) = delete;
+  ComputeBlockIdManager &operator=(ComputeBlockIdManager &&) = delete;
+
 private:
-  int cntComputeBlockId;
+  int cntComputeBlockId = 0;
   llvm::DenseMap<int, llvm::SmallVector<Operation *>> blockIdToOps;
   llvm::DenseMap<Operation *, int> opToBlockId;
-  mutable std::mutex managerMutex;
-  const int blockIdWidth = 32;
+  static constexpr int kBlockIdWidth = 32;
   llvm::LogicalResult markAndRecord(Operation *op, int blockId);
 };
 

third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h

@@ -26,11 +26,8 @@
 #include <memory>
 
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/Operation.h"
-#include "mlir/Pass/Pass.h"
 
-#include "DynamicCVPipeline/PlanComputeBlock/Common.h"
-#include "DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h"
+#include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace triton {
@@ -44,15 +41,6 @@ class PlanCubeBlockPass
   void runOnOperation() override;
 
   llvm::StringRef getArgument() const final { return "plan-cube-block"; }
-
-private:
-  SmallVector<Operation *>
-  matchSeed(Operation *dotOp, CVPipeline::ComputeBlockIdManager &bm,
-            const CVPipeline::MemoryDependenceGraph &memGraph);
-  llvm::LogicalResult
-  processBlockWithCubeBFS(Block *block,
-                          const CVPipeline::MemoryDependenceGraph &memGraph,
-                          CVPipeline::ComputeBlockIdManager &bm);
 };
 
 std::unique_ptr<OperationPass<ModuleOp>> createPlanCubeBlockPass();

third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp

@@ -40,7 +40,7 @@ using namespace triton;
 
 namespace {
 
-static constexpr llvm::StringLiteral interceptrFunc[]{""};
+static constexpr llvm::StringLiteral interceptrFunc[]{"_fwd_kernel_alibi"};
 
 static LogicalResult verifyFuncNames(ModuleOp module) {
   bool intercepted = false;

third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp

@@ -34,9 +34,11 @@
 // Unknown ops (no SideEffect interface) act as full barriers: they depend on
 // all prior writers/readers and become the sole writer for every slot.
 
-#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h"
-#include "ascend/include/DynamicCVPipeline/Common/Utils.h"
-#include "bishengir/Dialect/Annotation/IR/Annotation.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -47,10 +49,10 @@
 #include "mlir/IR/Region.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SetVector.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Debug.h"
+
+#include "DynamicCVPipeline/Common/MemoryEffectsTracker.h"
+#include "DynamicCVPipeline/Common/Utils.h"
+#include "bishengir/Dialect/Annotation/IR/Annotation.h"
 
 using namespace mlir;
 static constexpr const char *DEBUG_TYPE = "memory-effects-tracker";
@@ -332,12 +334,20 @@ MemoryDependenceGraph::collectOuterEffects(Operation *op, bool &unknown,
 }
 
 AliasResult MemoryDependenceGraph::queryAlias(Value lhs, Value rhs) {
+  auto lhsSource = getViewSource(lhs);
+  auto rhsSource = getViewSource(rhs);
+  if (!lhsSource) {
+    lhsSource = lhs;
+  }
+  if (!rhsSource) {
+    rhsSource = rhs;
+  }
+
   auto isFuncEntryArg = [](const Value &val) -> bool {
     auto arg = llvm::dyn_cast<BlockArgument>(val);
     return arg && arg.getOwner()->isEntryBlock();
   };
-  if (isFuncEntryArg(getViewSource(lhs)) &&
-      isFuncEntryArg(getViewSource(rhs))) {
+  if (isFuncEntryArg(lhsSource) && isFuncEntryArg(rhsSource)) {
     return lhs == rhs ? AliasResult::MustAlias : AliasResult::NoAlias;
   }
   return aa.alias(lhs, rhs);

third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp

@@ -2,9 +2,10 @@
 #include <optional>
 
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/LogicalResult.h"
 
-#include "bishengir/Dialect/HIVM/IR/HIVM.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -15,10 +16,20 @@
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 
 #include "ascend/include/DynamicCVPipeline/Common/Utils.h"
 
+#include "bishengir/Dialect/HIVM/IR/HIVM.h"
+
+static constexpr const char *DEBUG_TYPE = "dynamic-cv-pipeline-utils";
+#define DBGS(...) LLVM_DEBUG(llvm::dbgs() << __VA_ARGS__)
+#define LOG_DEBUG(...) DBGS("[" << DEBUG_TYPE << "] " << __VA_ARGS__)
+
 namespace mlir {
 namespace CVPipeline {
 
@@ -135,6 +146,55 @@ bool isOnlyDirectlyUse(Operation *preOp, Operation *nextOp,
   return (*allusers.begin()) == nextOp;
 }
 
+CoreType getCoreTypeOfSimpleOpOrCf(Operation *op) {
+  if (op == nullptr) {
+    return CoreType::UNDETERMINED;
+  }
+  if (!llvm::isa<RegionBranchOpInterface>(op)) {
+    return getOpCoreType(op);
+  }
+
+  CoreType coreType = CoreType::UNDETERMINED;
+  Operation *failingOp = nullptr;
+
+  // we need to skip sub-op of non-cf ops with regions, hence preorder here
+  op->walk<WalkOrder::PreOrder>([&](Operation *subOp) -> WalkResult {
+    if (llvm::isa<RegionBranchOpInterface>(subOp) ||
+        subOp->hasTrait<OpTrait::IsTerminator>()) {
+      return WalkResult::advance();
+    }
+
+    CoreType currCoreType = getOpCoreType(subOp);
+    // we have met a simple op without core type
+    if (currCoreType == CoreType::UNDETERMINED) {
+      coreType = CoreType::UNDETERMINED;
+      failingOp = subOp;
+      return WalkResult::interrupt();
+    }
+
+    if (coreType == CoreType::UNDETERMINED) {
+      coreType = currCoreType;
+    } else if (currCoreType != coreType) {
+      // some ops have different core type
+      coreType = CoreType::CUBE_AND_VECTOR;
+      return WalkResult::interrupt();
+    }
+
+    // skip sub-op
+    return WalkResult::skip();
+  });
+
+  (void)failingOp;
+  LOG_DEBUG("CoreType of RegionBranchOp is " << coreType << ": " << *op
+                                             << "\n");
+  LLVM_DEBUG({
+    if (coreType == CoreType::UNDETERMINED && failingOp != nullptr) {
+      llvm::dbgs() << "\nCoreType is UNDETERMINED due to " << *failingOp;
+    }
+  });
+  return coreType;
+}
+
 /** Determines if a value is "scalar-like" based on the following criteria:
  1. True scalar types (integer, index, or float)
  2. Tensor types with empty shape (e.g., tensor<f32>)
@@ -293,5 +353,37 @@ int64_t getBTSizeFromValidBroadcastOp(linalg::BroadcastOp broadcastOp) {
   return sizeBytes;
 }
 
+int getLoopCarriedArgIndex(Value operand, Block *block) {
+  if (!block || !block->mightHaveTerminator()) {
+    return -1;
+  }
+
+  auto barg = dyn_cast_if_present<BlockArgument>(operand);
+  if (!barg || barg.getOwner() != block) {
+    return -1;
+  }
+
+  auto *parentOp = block->getParentOp();
+  if (!isa<scf::ForOp, scf::WhileOp>(parentOp)) {
+    return -1;
+  }
+
+  auto *terminator = block->getTerminator();
+  if (!llvm::isa_and_present<scf::YieldOp>(terminator)) {
+    return -1;
+  }
+
+  int numArgs = block->getNumArguments();
+  int numYieldOperands = terminator->getNumOperands();
+  int offset = numArgs - numYieldOperands;
+  int argIdx = barg.getArgNumber() - offset;
+
+  if (argIdx < 0 || argIdx >= numYieldOperands) {
+    return -1;
+  }
+
+  return argIdx;
+}
+
 } // namespace CVPipeline
 } // namespace mlir