From 0c30fa28955c4b3519261167bcb443c1399ea071 Mon Sep 17 00:00:00 2001 From: chenwuyang Date: Thu, 16 Jul 2026 19:13:24 +0800 Subject: [PATCH 1/4] regional --- .../include/DynamicCVPipeline/Common/Utils.h | 52 ++- .../PlanComputeBlock/Common.h | 49 ++- .../PlanComputeBlock/ComputeBlockIdManager.h | 27 +- .../PlanComputeBlock/PlanCubeBlockPass.h | 14 +- .../AnalyzeDataFlow/AnalyzeName.cpp | 2 +- .../Common/MemoryEffectsTracker.cpp | 28 +- .../lib/DynamicCVPipeline/Common/Utils.cpp | 92 ++++- .../PlanComputeBlock/Common.cpp | 79 +++- .../ComputeBlockIdManager.cpp | 107 +++-- .../PlanComputeBlock/OpClassifier.cpp | 34 +- .../PlanComputeBlock/PlanCubeBlock.cpp | 262 +++++------- .../PlanComputeBlock/PlanVectorBlockPass.cpp | 383 ++++++++---------- .../PlanComputeBlock/ReorderOpsByBlockId.cpp | 27 +- .../SplitDataflow/RefineArgsBlockId.cpp | 18 +- 14 files changed, 669 insertions(+), 505 deletions(-) diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h index 1b39284f24..b41290cef0 100644 --- a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h +++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h @@ -86,6 +86,18 @@ constexpr int64_t BYTE_SIZE = 8; static constexpr int crossCoreProducerId = 1; static constexpr int crossCoreConsumerId = 0; +class MoveOnly { +protected: + MoveOnly() = default; + ~MoveOnly() = default; + + MoveOnly(const MoveOnly &) = delete; + MoveOnly &operator=(const MoveOnly &) = delete; + + MoveOnly(MoveOnly &&) = default; + MoveOnly &operator=(MoveOnly &&) = default; +}; + enum CoreType { UNDETERMINED = 0, VECTOR_ONLY = 1 << 0, @@ -121,8 +133,23 @@ 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; +} + +template +inline llvm::LogicalResult allSucceededShortCircuit(RangeT &&range, + FuncT &&func) { + return llvm::success( + llvm::all_of(std::forward(range), [&](auto &&item) { + return llvm::succeeded(func(std::forward(item))); + })); } bool isVectorOnlyOp(Operation *op); @@ -150,6 +177,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 diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h index 654a1eb981..4ce5b44ac3 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h +++ b/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; + + template + static auto mapToAncestorInBlock(Block *block, Fn &&pred) { + return [block, pred = std::forward(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; + + template + void forEachSource(Operation *op, PredFn pred) const; + + void forEachUserInSameBlock(Operation *op, PredFn pred) const { + forEachUser(op, mapToAncestorInBlock(op->getBlock(), pred)); + } + + template + void forEachSourceInSameBlock(Operation *op, PredFn pred) const { + forEachSource(op, + mapToAncestorInBlock(op->getBlock(), pred)); + } +}; + Operation *getAncestorInBlock(Operation *inner, Block *block); void initializeIndegreeForBlock(Block *block, llvm::DenseMap &indegree, - const MemoryDependenceGraph &memGraph, + const DependencyHelper &depHelper, ComputeBlockIdManager &bm); } // namespace CVPipeline diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h index 62e7c93e3a..4443c740f8 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h +++ b/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 + #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" #include "llvm/Support/LogicalResult.h" -#include + +#include "mlir/IR/Operation.h" + +#include "DynamicCVPipeline/Common/Utils.h" namespace mlir { namespace CVPipeline { @@ -36,7 +39,7 @@ namespace CVPipeline { /** * the class is to promise CUBEID and VECTORID are unified. */ -class ComputeBlockIdManager { +class ComputeBlockIdManager : MoveOnly { public: ComputeBlockIdManager(Operation *root); bool isSameBlock(Operation *a, Operation *b); @@ -47,17 +50,21 @@ class ComputeBlockIdManager { llvm::LogicalResult markOpsWithNewId(llvm::SmallVectorImpl &ops); void updateBlockId(Operation *op, int blockId); - llvm::SmallVector getOpsByBlockId(int blockId); - int getBlockIdByOp(Operation *op); - void reset(); + bool shouldInheritFromParent(Block *block, CoreType requiredCoreType) const; + llvm::LogicalResult inheritFromParent(Block *block); + + llvm::SmallVector getOpsByBlockId(int blockId) const; + llvm::SmallVector getOpsInSameBlock(Operation *op) const; + std::optional getBlockIdByOpOpt(Operation *op) const; int getNextId(); + int getBlockIdByOp(Operation *op); + private: - int cntComputeBlockId; + int cntComputeBlockId = 0; llvm::DenseMap> blockIdToOps; llvm::DenseMap opToBlockId; - mutable std::mutex managerMutex; - const int blockIdWidth = 32; + static constexpr int kBlockIdWidth = 32; llvm::LogicalResult markAndRecord(Operation *op, int blockId); }; diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h index 242e18c2ba..6bd028cddf 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h @@ -26,11 +26,8 @@ #include #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 - 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> createPlanCubeBlockPass(); diff --git a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp b/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp index 80e407255c..fb1ccc2ba0 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/AnalyzeDataFlow/AnalyzeName.cpp +++ b/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; diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp index a240134523..02f1b16173 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/Common/MemoryEffectsTracker.cpp +++ b/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(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); diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp index 37959eeb2f..283e04fadd 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp @@ -2,9 +2,10 @@ #include #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 = "DynamicCVPipeline/Utils"; +#define DBGS(...) LLVM_DEBUG(llvm::dbgs() << __VA_ARGS__) +#define LOG_DEBUG(...) DBGS("\n[" << DEBUG_TYPE << "] " << __VA_ARGS__) + namespace mlir { namespace CVPipeline { @@ -135,6 +146,53 @@ bool isOnlyDirectlyUse(Operation *preOp, Operation *nextOp, return (*allusers.begin()) == nextOp; } +CoreType getCoreTypeOfSimpleOpOrCf(Operation *op) { + if (op == nullptr) { + return CoreType::UNDETERMINED; + } + if (!llvm::isa(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([&](Operation *subOp) -> WalkResult { + if (llvm::isa(subOp) || + subOp->hasTrait()) { + 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(); + }); + + LOG_DEBUG("CoreType of RegionBranchOp is " << coreType << ": " << *op); + 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) @@ -293,5 +351,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(operand); + if (!barg || barg.getOwner() != block) { + return -1; + } + + auto *parentOp = block->getParentOp(); + if (!isa(parentOp)) { + return -1; + } + + auto *terminator = block->getTerminator(); + if (!llvm::isa_and_present(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 diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp index bd5b252b79..f13436836a 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp @@ -20,42 +20,79 @@ * THE SOFTWARE. */ +#include "llvm/ADT/iterator.h" + +#include "mlir/IR/Value.h" + #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" + +#include "DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "DynamicCVPipeline/Common/Utils.h" #include "DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" namespace mlir { namespace CVPipeline { -void initializeIndegreeForBlock(Block *block, - llvm::DenseMap &indegree, - const MemoryDependenceGraph &memGraph, - ComputeBlockIdManager &bm) { +void DependencyHelper::forEachUser(Operation *op, + DependencyHelper::PredFn pred) const { + for (auto *user : op->getUsers()) { + pred(user); + } + for (auto *user : memGraph.getExecAfter(op)) { + pred(user); + } +} - block->walk([&](Operation *op) { - if (op->getBlock() != block) { - return; - } - indegree[op] = 0; - // We need to consider op itself && op's region-contained ops. - op->walk([&](Operation *nestedOp) { - for (auto inValue : nestedOp->getOperands()) { - if (auto defOp = inValue.getDefiningOp()) { - if (defOp->getBlock() == block && !bm.isSameBlock(defOp, op)) { - indegree[op]++; - } +template +void DependencyHelper::forEachSource(Operation *op, + DependencyHelper::PredFn pred) const { + op->walk([&, this, op](Operation *subOp) { + for (auto operand : subOp->getOperands()) { + if (auto *defOp = operand.getDefiningOp(); defOp) { + if (!op->isAncestor(defOp)) { + pred(defOp); } + continue; } - for (auto memDepUser : memGraph.getExecBefore(nestedOp)) { - if (memDepUser->getBlock() == block && - !bm.isSameBlock(memDepUser, op)) { - indegree[op]++; + if constexpr (AcrossIterArg) { + if (auto *defOp = getLoopCarriedDefOp(operand, op->getBlock())) { + pred(defOp); } } - }); + } + for (auto *source : memGraph.getExecBefore(subOp)) { + if (!op->isAncestor( + source)) { // this filters only the outer mem dependencies + pred(source); + } + } }); } +// Instantiate concrete functions for linking +template void mlir::CVPipeline::DependencyHelper::forEachSource( + mlir::Operation *op, + llvm::function_ref callback) const; + +template void mlir::CVPipeline::DependencyHelper::forEachSource( + mlir::Operation *op, + llvm::function_ref callback) const; + +void initializeIndegreeForBlock(Block *block, + llvm::DenseMap &indegree, + const DependencyHelper &depHelper, + ComputeBlockIdManager &bm) { + for (auto *op : llvm::make_pointer_range(block->getOperations())) { + indegree[op] = 0; + depHelper.forEachSource(op, [&](Operation *source) { + if (source->getBlock() == block && !bm.isSameBlock(source, op)) { + indegree[op]++; + } + }); + } +} + Operation *getAncestorInBlock(Operation *inner, Block *block) { Operation *cur = inner; while (cur) { diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp index 1d293b3842..fb5b2d331c 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp @@ -20,29 +20,33 @@ * THE SOFTWARE. */ -#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" -#include "ascend/include/DynamicCVPipeline/Common/Utils.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinTypes.h" +#include + +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Debug.h" #include "llvm/Support/LogicalResult.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" + namespace mlir { namespace CVPipeline { ComputeBlockIdManager::ComputeBlockIdManager(Operation *root) { - cntComputeBlockId = 0; blockIdToOps.clear(); opToBlockId.clear(); root->walk([&](Operation *op) { if (auto blockIdAttr = op->getAttrOfType(kBlockId)) { - if (auto blockId = blockIdAttr.getInt()) { - opToBlockId[op] = blockId; - blockIdToOps[blockId].push_back(op); - cntComputeBlockId = - std::max(cntComputeBlockId, static_cast(blockId)); - } + auto blockId = blockIdAttr.getInt(); + opToBlockId[op] = blockId; + blockIdToOps[blockId].push_back(op); + cntComputeBlockId = + std::max(cntComputeBlockId, static_cast(blockId)); } }); cntComputeBlockId++; // ensure new id is unique @@ -50,12 +54,7 @@ ComputeBlockIdManager::ComputeBlockIdManager(Operation *root) { bool ComputeBlockIdManager::isWholeCubeReady( Operation *seedOp, llvm::DenseMap &indegree) { - auto id = getBlockIdByOp(seedOp); - if (id == -1) { - return (indegree[seedOp] == 0); - } - auto cubeBlock = getOpsByBlockId(id); - for (auto op : cubeBlock) { + for (auto *op : getOpsInSameBlock(seedOp)) { if (!indegree.contains(op)) { continue; } @@ -84,9 +83,10 @@ void ComputeBlockIdManager::updateBlockId(Operation *op, int blockId) { if (blockId == -1) { op->removeAttr(kBlockId); } else { - op->setAttr(kBlockId, - IntegerAttr::get(IntegerType::get(ctx, blockIdWidth), blockId)); + op->setAttr(kBlockId, IntegerAttr::get(IntegerType::get(ctx, kBlockIdWidth), + blockId)); } + auto it = opToBlockId.find(op); if (it != opToBlockId.end()) { int preBlockId = it->second; @@ -98,12 +98,13 @@ void ComputeBlockIdManager::updateBlockId(Operation *op, int blockId) { } } } + opToBlockId[op] = blockId; blockIdToOps[blockId].push_back(op); } llvm::SmallVector -ComputeBlockIdManager::getOpsByBlockId(int blockId) { +ComputeBlockIdManager::getOpsByBlockId(int blockId) const { if (blockId == -1) { return {}; } @@ -115,12 +116,35 @@ ComputeBlockIdManager::getOpsByBlockId(int blockId) { return llvm::SmallVector(it->second.begin(), it->second.end()); } -int ComputeBlockIdManager::getBlockIdByOp(Operation *op) { +llvm::SmallVector +ComputeBlockIdManager::getOpsInSameBlock(Operation *op) const { + auto blockIdOpt = getBlockIdByOpOpt(op); + if (!blockIdOpt.has_value()) { + return {op}; + } + auto blockId = blockIdOpt.value(); + auto *block = op->getBlock(); + if (auto it = blockIdToOps.find(blockId); it != blockIdToOps.end()) { + auto filtered = + llvm::make_filter_range(it->second, [block](Operation *opInBlock) { + return opInBlock->getBlock() == block; + }); + return {filtered.begin(), filtered.end()}; + } + return {op}; +} + +std::optional +ComputeBlockIdManager::getBlockIdByOpOpt(Operation *op) const { auto it = opToBlockId.find(op); if (it != opToBlockId.end()) { return it->second; } - return -1; + return std::nullopt; +} + +int ComputeBlockIdManager::getBlockIdByOp(Operation *op) { + return getBlockIdByOpOpt(op).value_or(-1); } llvm::LogicalResult ComputeBlockIdManager::markAndRecord(Operation *op, @@ -128,7 +152,7 @@ llvm::LogicalResult ComputeBlockIdManager::markAndRecord(Operation *op, // When we call mark, we assume the op have no record in manager. MLIRContext *ctx = op->getContext(); op->setAttr(kBlockId, - IntegerAttr::get(IntegerType::get(ctx, blockIdWidth), blockId)); + IntegerAttr::get(IntegerType::get(ctx, kBlockIdWidth), blockId)); auto itOld = opToBlockId.find(op); if (itOld != opToBlockId.end() && itOld->second != -1) { llvm::errs() << "Error: Operation already has a block id. Op: " << *op @@ -162,10 +186,39 @@ llvm::LogicalResult ComputeBlockIdManager::markOpsWithNewId( return llvm::success(); } -void ComputeBlockIdManager::reset() { - cntComputeBlockId = 0; - blockIdToOps.clear(); - opToBlockId.clear(); +bool ComputeBlockIdManager::shouldInheritFromParent( + Block *block, CoreType requiredCoreType) const { + auto *parentOp = block->getParentOp(); + if (!parentOp || !isScfOp(parentOp) || + getCoreTypeOfSimpleOpOrCf(parentOp) != requiredCoreType) { + return false; + } + + auto blockIdOpt = getBlockIdByOpOpt(parentOp); + return blockIdOpt.has_value(); +} + +llvm::LogicalResult ComputeBlockIdManager::inheritFromParent(Block *block) { + auto *parentOp = block->getParentOp(); + if (!parentOp) { + return llvm::failure(); + } + + auto blockIdOpt = getBlockIdByOpOpt(parentOp); + if (!blockIdOpt.has_value()) { + return llvm::failure(); + } + + // no need to and should not walk inside nested blocks: + // 1. the caller is from walk already + // 2. if we have marked nested ops with this block id, it could be correct, + // since they will be re-marked with the same id, so no failures will be + // returned, but this is less robust + return allSucceededShortCircuit( + block->getOperations(), + [this, blockId = blockIdOpt.value()](Operation &op) { + return markAndRecord(&op, blockId); + }); } } // namespace CVPipeline diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp index f2043fb625..b9a3739d08 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp @@ -23,6 +23,7 @@ #include #include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" @@ -41,6 +42,8 @@ #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/OpClassifier.h" #include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" #include "bishengir/Dialect/Utils/Util.h" using namespace mlir; @@ -179,18 +182,25 @@ void OpClassifierPass::matchToTensorPattern(Operation *def) { if (!toTensorOp) return; - // special case: implicit transpose + // special case: implicit transpose -> vector if (utils::getAnnotateOpWithAttr(toTensorOp.getResult(), kMayImplicitTransposeWithLastAxis)) { return; } + Value memref = toTensorOp.getBuffer(); + // special case: ExtractLoadStore -> vector + if (llvm::any_of(memref.getUsers(), [](Operation *user) { + auto forOp = user->getParentOfType(); + return forOp && forOp->hasAttr(hivm::ExtractLoadStoreAttr); + })) { + return; + } + markCube(toTensorOp); cubeSeeds.push_back(toTensorOp); // Also mark the memref allocation as CUBE - Value memref = toTensorOp.getBuffer(); - if (Operation *memrefDef = memref.getDefiningOp()) { markCube(memrefDef); cubeSeeds.push_back(memrefDef); @@ -793,6 +803,24 @@ int OpClassifierPass::markRemainingAsVector() { if (opCoreTypes[op] == OP_UNDETERMINED && !isa(op)) { opCoreTypes[op] = OP_VECTOR_ONLY; } + + // ExtractLoadStoreAttr -> force on vector + if (isa(op) && op->hasAttr(hivm::ExtractLoadStoreAttr)) { + op->walk([this](Operation *nestedOp) { + opCoreTypes[nestedOp] = OP_VECTOR_ONLY; + for (auto operand : nestedOp->getOperands()) { + if (auto allocOp = llvm::dyn_cast_if_present( + operand.getDefiningOp())) { + opCoreTypes[allocOp] = OP_VECTOR_ONLY; + for (auto *user : allocOp->getUsers()) { + if (llvm::isa(user)) { + opCoreTypes[user] = OP_VECTOR_ONLY; + } + } + } + } + }); + } } return 0; diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp index de6518ea41..ad55d42e75 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp @@ -20,7 +20,6 @@ * THE SOFTWARE. */ -#include #include #include @@ -34,12 +33,9 @@ #include "llvm/Support/LogicalResult.h" #include "llvm/Support/raw_ostream.h" -#include "DynamicCVPipeline/ComputeBlockOpt/Common.h" #include "mlir/Analysis/AliasAnalysis.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Block.h" #include "mlir/IR/BuiltinOps.h" @@ -47,12 +43,14 @@ #include "mlir/IR/Visitors.h" #include "mlir/Pass/Pass.h" +#include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/ComputeBlockOpt/Common.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlockPass.h" -#include "DynamicCVPipeline/Common/MemoryEffectsTracker.h" -#include "DynamicCVPipeline/Common/Utils.h" #include "bishengir/Dialect/Annotation/IR/Annotation.h" -#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" using namespace mlir; using namespace triton; @@ -69,7 +67,7 @@ namespace { class SeedRegionPlanner { SmallVector seeds; Block *block; - const MemoryDependenceGraph &memGraph; + const DependencyHelper &depHelper; ComputeBlockIdManager &bm; llvm::DenseSet &assigned; llvm::SmallVectorImpl &group; @@ -79,11 +77,11 @@ class SeedRegionPlanner { public: SeedRegionPlanner(SmallVector seeds, Block *block, - const MemoryDependenceGraph &memGraph, + const DependencyHelper &depHelper, llvm::DenseSet &assigned, llvm::SmallVectorImpl &group, ComputeBlockIdManager &bm) - : seeds(seeds), block(block), memGraph(memGraph), assigned(assigned), + : seeds(seeds), block(block), depHelper(depHelper), assigned(assigned), group(group), bm(bm) { for (auto sd : seeds) { group.push_back(sd); @@ -100,17 +98,17 @@ namespace { class DependencyCycleDetector { const llvm::DenseSet &group; llvm::DenseSet visited; - const MemoryDependenceGraph &memGraph; + const DependencyHelper &depHelper; ComputeBlockIdManager &bm; Block *const block; bool detectCycleFrom(Operation *cur); public: - DependencyCycleDetector(Block *block, const MemoryDependenceGraph &memGraph, + DependencyCycleDetector(Block *block, const DependencyHelper &depHelper, llvm::DenseSet &group, ComputeBlockIdManager &bm) - : block(block), memGraph(memGraph), group(group), bm(bm) {} + : block(block), depHelper(depHelper), group(group), bm(bm) {} bool detectCycle(); }; @@ -125,42 +123,25 @@ bool DependencyCycleDetector::detectCycleFrom(Operation *cur) { return false; } - auto userCreatesCycle = [this, cur](Operation *user) { - auto *userInBlock = getAncestorInBlock(user, block); - if (!userInBlock) { - return false; - } - auto userBlockId = bm.getBlockIdByOp(userInBlock); - if (userBlockId == -1) { - return detectCycleFrom(userInBlock); - } - - return llvm::any_of( - bm.getOpsByBlockId(userBlockId), - [this](Operation *user) { return detectCycleFrom(user); }); - }; + bool createsCycle = false; - return llvm::any_of(cur->getUsers(), userCreatesCycle) || - llvm::any_of(memGraph.getExecAfter(cur), userCreatesCycle); -} + depHelper.forEachUserInSameBlock(cur, [&](Operation *user) { + createsCycle = createsCycle || llvm::any_of(bm.getOpsInSameBlock(user), + [this](Operation *user) { + return detectCycleFrom(user); + }); + return; + }); -static void forEachUser(Operation *op, const MemoryDependenceGraph &memGraph, - const std::function &pred) { - for (auto *user : op->getUsers()) { - pred(user); - } - for (auto *user : memGraph.getExecAfter(op)) { - pred(user); - } + return createsCycle; } bool DependencyCycleDetector::detectCycle() { llvm::DenseSet externalUsers; for (auto *op : group) { - forEachUser(op, memGraph, [&](Operation *user) { - auto *userInBlock = getAncestorInBlock(user, block); - if (userInBlock && !group.contains(userInBlock)) { - externalUsers.insert(userInBlock); + depHelper.forEachUserInSameBlock(op, [&](Operation *user) { + if (!group.contains(user)) { + externalUsers.insert(user); } }); } @@ -174,7 +155,7 @@ bool SeedRegionPlanner::willCreateCycle(Operation *op) { llvm::DenseSet okSet(group.begin(), group.end()); okSet.insert(op); - DependencyCycleDetector dfs = {block, memGraph, okSet, bm}; + DependencyCycleDetector dfs = {block, depHelper, okSet, bm}; return dfs.detectCycle(); } @@ -184,7 +165,7 @@ bool SeedRegionPlanner::willCreateCycle(Operation *op) { * op, and not creating a cycle in the dependence graph. */ bool SeedRegionPlanner::isEligible(Operation *op) { - if (!isCubeOp(op) || assigned.contains(op) || isMatmulOp(op)) { + if (!isCubeSimpleOpOrCf(op) || assigned.contains(op) || isMatmulOp(op)) { return false; } return !willCreateCycle(op); @@ -203,29 +184,8 @@ void SeedRegionPlanner::run() { size_t head = 0; while (head < group.size()) { Operation *currOp = group[head++]; - - // Check data operands - for (Value iop : currOp->getOperands()) { - if (auto *def = iop.getDefiningOp()) { - tryAddToGroup(def); - } - // Check loop-carried dependencies (SCF ForOp iter_args) - if (auto barg = dyn_cast(iop)) { - if (barg.getOwner() == block && isa(block->getParentOp()) && - barg.getArgNumber() > 0) { - auto *yieldOp = barg.getOwner()->getTerminator(); - if (auto *yieldedValDef = yieldOp->getOperand(barg.getArgNumber() - 1) - .getDefiningOp()) { - tryAddToGroup(yieldedValDef); - } - } - } - } - - // Check memory dependencies (RAW/WAW/WAR) - for (auto *def : memGraph.getMemDefs(currOp)) { - tryAddToGroup(def); - } + depHelper.forEachSource( + currOp, [this](Operation *source) { tryAddToGroup(source); }); } } @@ -240,7 +200,7 @@ class TopologicalPartitionPlanner { unsigned nonAssignedCubeCnt = 0; llvm::DenseMap indegree; llvm::DenseSet &assigned; - const MemoryDependenceGraph &memGraph; + const DependencyHelper &depHelper; ComputeBlockIdManager &bm; llvm::DenseSet newassigned; llvm::DenseSet bypassVisited; @@ -248,8 +208,9 @@ class TopologicalPartitionPlanner { void removeNonCubeOpsRecursively(Operation *op); llvm::LogicalResult removeReadyNonCubeOps(); + bool shouldSkip(Operation *op) { - return !isCubeOp(op) || assigned.contains(op); + return !isCubeSimpleOpOrCf(op) || assigned.contains(op); }; bool canExpandTo(Operation *op); void dumpQueueAndIndegreeInfo(); @@ -259,18 +220,18 @@ class TopologicalPartitionPlanner { public: TopologicalPartitionPlanner(Block *block, llvm::DenseSet &assigned, - const MemoryDependenceGraph &memGraph, + const DependencyHelper &depHelper, ComputeBlockIdManager &bm) - : block(block), assigned(assigned), memGraph(memGraph), bm(bm) { - initializeIndegreeForBlock(block, indegree, memGraph, bm); + : block(block), assigned(assigned), depHelper(depHelper), bm(bm) { + initializeIndegreeForBlock(block, indegree, depHelper, bm); block->walk([&](Operation *op) { - if (op->getBlock() == block && isCubeOp(op) && !assigned.contains(op)) { + if (op->getBlock() == block && isCubeSimpleOpOrCf(op) && + !assigned.contains(op)) { nonAssignedCubeCnt++; } }); } - llvm::LogicalResult run(); }; @@ -281,47 +242,22 @@ class TopologicalPartitionPlanner { void TopologicalPartitionPlanner::removeNonCubeOpsRecursively(Operation *op) { LOG_DEBUG("\tRemoved non-cube:" << *op << "\n"); bypassVisited.insert(op); - auto *block = op->getBlock(); - SmallVector allusers; - allusers.append(op->getUsers().begin(), op->getUsers().end()); - for (auto *memUser : memGraph.getExecAfter(op)) { - allusers.push_back(memUser); - } - for (auto *user : allusers) { - auto *userInBlock = getAncestorInBlock(user, block); - if (!userInBlock || !indegree.contains(userInBlock) || - bm.isSameBlock(userInBlock, op)) { - continue; - } - LOG_DEBUG("Sub indegree to " - << *userInBlock << " from " << *op - << "new degree = " << indegree[userInBlock] - 1 << "\n"); - indegree[userInBlock]--; - if (!bm.isWholeCubeReady(userInBlock, indegree) || - bypassVisited.contains(userInBlock) || !shouldSkip(userInBlock)) { - continue; - } - auto blockId = bm.getBlockIdByOp(userInBlock); - if (blockId == -1) { - removeNonCubeOpsRecursively(userInBlock); - continue; - } - for (auto *passop : bm.getOpsByBlockId(blockId)) { + depHelper.forEachUserInSameBlock(op, [&](Operation *user) { + if (!indegree.contains(user) || bm.isSameBlock(user, op)) { + return; + } + LOG_DEBUG("Sub indegree to " << *user << " from " << *op << "new degree = " + << indegree[user] - 1 << "\n"); + indegree[user]--; + if (!bm.isWholeCubeReady(user, indegree) || bypassVisited.contains(user) || + !shouldSkip(user)) { + return; + } + for (auto *passop : bm.getOpsInSameBlock(user)) { if (!bypassVisited.contains(passop)) { removeNonCubeOpsRecursively(passop); } } - } -} - -static bool mapsAreDiff(const llvm::DenseMap &a, - const llvm::DenseMap &b) { - if (a.size() != b.size()) { - return true; - } - return llvm::any_of(a, [&b](std::pair aIter) { - auto bIter = b.find(aIter.first); - return bIter == b.end() || bIter->second != aIter.second; }); } @@ -334,22 +270,17 @@ llvm::LogicalResult TopologicalPartitionPlanner::removeReadyNonCubeOps() { size_t beforeVisitedSize = bypassVisited.size(); for (auto &p : indegree) { Operation *op = p.first; - if (shouldSkip(op) && bm.isWholeCubeReady(op, indegree) && - !bypassVisited.contains(op)) { - int blockId = bm.getBlockIdByOp(op); - if (blockId == -1) { - removeNonCubeOpsRecursively(op); - } else { - for (auto *passOp : bm.getOpsByBlockId(blockId)) { - if (!bypassVisited.contains(passOp)) { - removeNonCubeOpsRecursively(passOp); - } - } + if (!shouldSkip(op) || !bm.isWholeCubeReady(op, indegree) || + bypassVisited.contains(op)) { + continue; + } + for (auto *passOp : bm.getOpsInSameBlock(op)) { + if (!bypassVisited.contains(passOp)) { + removeNonCubeOpsRecursively(passOp); } } } - if (!mapsAreDiff(indegreeBefore, indegree) && - beforeVisitedSize == bypassVisited.size()) { + if (indegreeBefore == indegree && beforeVisitedSize == bypassVisited.size()) { if (Operation *parentOp = block->getParentOp()) { parentOp->emitError("PlanCubeBlock cannot make progress while scheduling " "cube operations"); @@ -363,7 +294,7 @@ llvm::LogicalResult TopologicalPartitionPlanner::removeReadyNonCubeOps() { // Expansion condition: op must be CUBE_ONLY, indegree == 0 and all its // dependency ops are CUBE_ONLY bool TopologicalPartitionPlanner::canExpandTo(Operation *op) { - if (!isCubeOp(op) || assigned.contains(op)) { + if (!isCubeSimpleOpOrCf(op) || assigned.contains(op)) { return false; } auto it = indegree.find(op); @@ -396,7 +327,7 @@ void TopologicalPartitionPlanner::dumpQueueAndIndegreeInfo() { bool foundRemainingCube = false; for (auto &p : indegree) { Operation *op = p.first; - if (!op || op->getBlock() != block || !CVPipeline::isCubeOp(op) || + if (!op || op->getBlock() != block || !isCubeSimpleOpOrCf(op) || assigned.contains(op) || newassigned.contains(op)) { continue; } @@ -414,7 +345,7 @@ llvm::LogicalResult TopologicalPartitionPlanner::populateQueueWithReadyOps() { op->emitError("Indegree cannot be negative"); return llvm::failure(); } - if (indegree == 0 && !newassigned.contains(op) && isCubeOp(op) && + if (indegree == 0 && !newassigned.contains(op) && isCubeSimpleOpOrCf(op) && !assigned.contains(op)) { queue.push(op); } @@ -433,24 +364,18 @@ TopologicalPartitionPlanner::createNewGroupFromQueue() { group.push_back(currOp); nonAssignedCubeCnt--; - llvm::SmallVector allUsers; - for (auto *user : currOp->getUsers()) - allUsers.push_back(user); - for (auto *user : memGraph.getExecAfter(currOp)) - allUsers.push_back(user); - for (auto *user : allUsers) { - auto *userInBlock = getAncestorInBlock(user, block); - if (userInBlock && !newassigned.contains(userInBlock)) { - auto &userInDegree = indegree[userInBlock]; + depHelper.forEachUserInSameBlock(currOp, [&](Operation *user) { + if (!newassigned.contains(user)) { + auto &userInDegree = indegree[user]; userInDegree--; - LOG_DEBUG("Sub indegree to " << *userInBlock << " from " << *currOp + LOG_DEBUG("Sub indegree to " << *user << " from " << *currOp << "new degree = " << userInDegree << "\n"); - if (canExpandTo(userInBlock)) { - queue.push(userInBlock); + if (canExpandTo(user)) { + queue.push(user); } } - } + }); } return group; } @@ -488,7 +413,7 @@ static SmallVector collectMatmulOps(Block *block) { } static void fuseMarkOpToDef(Block *block, ComputeBlockIdManager &bm, - const MemoryDependenceGraph &memGraph) { + const DependencyHelper &depHelper) { for (auto *op : llvm::make_pointer_range(block->getOperations())) { if (getOpCoreType(op) != CUBE_ONLY) { continue; @@ -507,7 +432,7 @@ static void fuseMarkOpToDef(Block *block, ComputeBlockIdManager &bm, continue; } - auto currGroup = bm.getOpsByBlockId(defBlockId); + auto currGroup = bm.getOpsInSameBlock(defOp); llvm::DenseSet newGroup{currGroup.begin(), currGroup.end()}; if (newGroup.contains(markOp)) { @@ -515,7 +440,7 @@ static void fuseMarkOpToDef(Block *block, ComputeBlockIdManager &bm, } newGroup.insert(markOp); - DependencyCycleDetector dfs{block, memGraph, newGroup, bm}; + DependencyCycleDetector dfs{block, depHelper, newGroup, bm}; if (!dfs.detectCycle()) { bm.updateBlockId(markOp, defBlockId); } @@ -533,9 +458,10 @@ static bool checkValidUserSeed(Operation *op) { return isa(op); } -SmallVector -PlanCubeBlockPass::matchSeed(Operation *dotOp, ComputeBlockIdManager &bm, - const MemoryDependenceGraph &memGraph) { + +static SmallVector +matchSeed(Operation *dotOp, ComputeBlockIdManager &bm, + const MemoryDependenceGraph &memGraph) { // match inputs SmallVector ret; ret.push_back(dotOp); @@ -543,7 +469,7 @@ PlanCubeBlockPass::matchSeed(Operation *dotOp, ComputeBlockIdManager &bm, Operation *def = operand.getDefiningOp(); if (!def) continue; - if (checkValidInputSeed(def) && isCubeOp(def) && + if (checkValidInputSeed(def) && isCubeSimpleOpOrCf(def) && dotOp->getBlock() == def->getBlock() && bm.getBlockIdByOp(def) == -1) { if (CVPipeline::isOnlyDirectlyUse(def, dotOp, memGraph)) { ret.push_back(def); @@ -554,7 +480,7 @@ PlanCubeBlockPass::matchSeed(Operation *dotOp, ComputeBlockIdManager &bm, Operation *nowOp = dotOp; while (nowOp->hasOneUse()) { auto user = *nowOp->getUsers().begin(); - if (user->getBlock() != dotOp->getBlock() || !isCubeOp(user) || + if (user->getBlock() != dotOp->getBlock() || !isCubeSimpleOpOrCf(user) || bm.getBlockIdByOp(user) != -1) { break; } @@ -572,9 +498,9 @@ PlanCubeBlockPass::matchSeed(Operation *dotOp, ComputeBlockIdManager &bm, * Main entry point: Process a single block by grouping operations into * execution blocks using BFS and topological traversal. */ -llvm::LogicalResult PlanCubeBlockPass::processBlockWithCubeBFS( - Block *block, const MemoryDependenceGraph &memGraph, - ComputeBlockIdManager &bm) { +static llvm::LogicalResult +processBlockWithCubeBFS(Block *block, const DependencyHelper &depHelper, + ComputeBlockIdManager &bm) { llvm::DenseSet assigned; auto allDots = collectMatmulOps(block); @@ -585,8 +511,9 @@ llvm::LogicalResult PlanCubeBlockPass::processBlockWithCubeBFS( continue; } auto temBlockId = bm.getNextId(); - llvm::SmallVector dotSeeds = matchSeed(dot, bm, memGraph); - if (willCreateCycle(dotSeeds, memGraph, temBlockId, bm)) { + llvm::SmallVector dotSeeds = + matchSeed(dot, bm, depHelper.memGraph); + if (willCreateCycle(dotSeeds, depHelper.memGraph, temBlockId, bm)) { LOG_DEBUG("Cube Seed already have a cycle!!"); for (auto seed : dotSeeds) { LOG_DEBUG("Seed: " << *seed << "\n"); @@ -594,7 +521,7 @@ llvm::LogicalResult PlanCubeBlockPass::processBlockWithCubeBFS( return llvm::failure(); } llvm::SmallVector newGroup; - SeedRegionPlanner regionPlanner{dotSeeds, block, memGraph, + SeedRegionPlanner regionPlanner{dotSeeds, block, depHelper, assigned, newGroup, bm}; regionPlanner.run(); @@ -607,31 +534,42 @@ llvm::LogicalResult PlanCubeBlockPass::processBlockWithCubeBFS( } // Phase 2: Handle remaining Cube Ops following Topo order - TopologicalPartitionPlanner topoPlanner{block, assigned, memGraph, bm}; + TopologicalPartitionPlanner topoPlanner{block, assigned, depHelper, bm}; if (failed(topoPlanner.run())) { return failure(); } - fuseMarkOpToDef(block, bm, memGraph); + fuseMarkOpToDef(block, bm, depHelper); return llvm::success(); } void mlir::triton::PlanCubeBlockPass::runOnOperation() { - LOG_DEBUG( - "\n--- Step 2: Partitioning compute blocks for cube operations --->\n"); auto moduleOp = getOperation(); if (CVPipeline::hasFallbackAttr(moduleOp)) { return; } + LOG_DEBUG("Input mlir:\n" << moduleOp << "\n==========\n"); + auto &aa = getAnalysis(); - auto memGraph = MemoryDependenceGraph(moduleOp, aa); + MemoryDependenceGraph memGraph{moduleOp, aa}; + DependencyHelper depHelper{memGraph}; auto bm = ComputeBlockIdManager(moduleOp); // We do not need to skip linalg blocks since they do not have core types and // do not contain matmul - auto result = moduleOp.walk([&](Block *block) { - if (llvm::failed(processBlockWithCubeBFS(block, memGraph, bm))) { + auto result = moduleOp.walk([&](Block *block) { + if (bm.shouldInheritFromParent(block, CoreType::CUBE_ONLY)) { + if (llvm::failed(bm.inheritFromParent(block))) { + block->getParentOp()->emitError() + << "[" << DEBUG_TYPE + << "] Sub-blocks failed to inherit block id from parent op"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + + if (llvm::failed(processBlockWithCubeBFS(block, depHelper, bm))) { return WalkResult::interrupt(); } return WalkResult::advance(); @@ -639,7 +577,7 @@ void mlir::triton::PlanCubeBlockPass::runOnOperation() { if (result.wasInterrupted()) { CVPipeline::setFallbackAttr(moduleOp, CVPipeline::ERRCODE_FAILED); } - LOG_DEBUG("\n--- Step 2: end --->\n"); + LOG_DEBUG("Output mlir:\n" << moduleOp << "\n==========\n"); } std::unique_ptr> diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp index 20f38cea90..a711d64da5 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp @@ -20,27 +20,32 @@ * THE SOFTWARE. */ -#include "DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" -#include "ascend/include/DynamicCVPipeline/Common/Utils.h" -#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" -#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h" -#include "bishengir/Dialect/Annotation/IR/Annotation.h" -#include "mlir/Analysis/AliasAnalysis.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/Block.h" -#include "mlir/IR/Operation.h" -#include "mlir/Support/LLVM.h" +#include + #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SetVector.h" -#include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" #include "llvm/Support/Debug.h" #include "llvm/Support/LogicalResult.h" #include "llvm/Support/raw_ostream.h" -#include -#include + +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Interfaces/ControlFlowInterfaces.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/WalkResult.h" + +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Passes.h" + +#include "DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" static constexpr const char *DEBUG_TYPE = "plan-vector-block"; #define LOG_DEBUG(...) \ @@ -50,25 +55,9 @@ using namespace mlir; using namespace triton; using namespace CVPipeline; -namespace mlir { -namespace triton { -class PlanVectorBlockPass - : public PassWrapper> { -public: - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PlanVectorBlockPass) - - PlanVectorBlockPass() = default; - void runOnOperation() override; - - llvm::StringRef getArgument() const final { return "plan-vector-block"; } -}; - -bool isFusableOp(Operation *op) { - if (CVPipeline::getOpCoreType(op) == CVPipeline::CoreType::VECTOR_ONLY) { - if (isa(op)) { - // pass control ops like scf::ForOp/scf::IfOp/scf::WhileOp - return false; - } +static bool isFusableOp(Operation *op) { + if (isVectorSimpleOpOrCf(op)) { + // skip terminators if (op->getBlock()->mightHaveTerminator() && op == op->getBlock()->getTerminator()) { return false; @@ -78,86 +67,71 @@ bool isFusableOp(Operation *op) { return false; } -void passAndCollectCandidates(Operation *nowOp, - DenseMap &indegree, - SmallVector &candidates, - DenseMap &visited, - const CVPipeline::MemoryDependenceGraph &memGraph, - ComputeBlockIdManager &bm) { +static void +passAndCollectCandidates(Operation *nowOp, DenseMap &indegree, + SmallVector &candidates, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph, + ComputeBlockIdManager &bm) { LOG_DEBUG("Bypassing non-fusable op " << *nowOp << "\nnow candidates size: " << candidates.size() << "\n"); - auto block = nowOp->getBlock(); - SmallVector allusers; - allusers.append(nowOp->getUsers().begin(), nowOp->getUsers().end()); - for (auto memUser : memGraph.getExecAfter(nowOp)) { - allusers.push_back(memUser); - } - for (auto user : allusers) { - auto userInBlock = CVPipeline::getAncestorInBlock(user, block); - if (!userInBlock) { - continue; + DependencyHelper depHelper{memGraph}; + depHelper.forEachUserInSameBlock(nowOp, [&](Operation *user) { + if (!bm.isSameBlock(user, nowOp)) { + indegree[user]--; } - if (!bm.isSameBlock(userInBlock, nowOp)) { - indegree[userInBlock]--; + + if (!bm.isWholeCubeReady(user, indegree)) { + return; } - if (bm.isWholeCubeReady(userInBlock, indegree)) { - if (!isFusableOp(userInBlock) && !visited[userInBlock]) { - if (bm.getBlockIdByOp(userInBlock) == -1) { - visited[userInBlock] = - true; // mark as fused to avoid duplicate bypass - passAndCollectCandidates(userInBlock, indegree, candidates, visited, - memGraph, bm); - } else { - for (auto cubeop : - bm.getOpsByBlockId(bm.getBlockIdByOp(userInBlock))) { - if (!visited[cubeop]) { - visited[cubeop] = true; - passAndCollectCandidates(cubeop, indegree, candidates, visited, - memGraph, bm); - } - } - } - } else if (isFusableOp(userInBlock) && !visited[userInBlock]) { - visited[userInBlock] = true; - candidates.push_back(userInBlock); + if (visited[user]) { + return; + } + + if (isFusableOp(user)) { + visited[user] = true; + candidates.push_back(user); + return; + } + + for (auto *cubeop : bm.getOpsInSameBlock(user)) { + if (!visited[cubeop]) { + visited[cubeop] = true; + passAndCollectCandidates(cubeop, indegree, candidates, visited, + memGraph, bm); } } - } + }); LOG_DEBUG("After bypassing, candidates size: " << candidates.size() << "\n"); } -void byPassNonFusable(DenseMap &indegree, - SmallVector &candidates, - DenseMap &visited, - const CVPipeline::MemoryDependenceGraph &memGraph, - ComputeBlockIdManager &bm) { +static void byPassNonFusable(DenseMap &indegree, + SmallVector &candidates, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph, + ComputeBlockIdManager &bm) { // for every non-fusable candidates, bypass it. - for (auto &[op, degree] : indegree) { - if (bm.isWholeCubeReady(op, indegree) && !isFusableOp(op) && !visited[op]) { - if (bm.getBlockIdByOp(op) == -1) { - visited[op] = true; // mark as fused to avoid duplicate bypass - passAndCollectCandidates(op, indegree, candidates, visited, memGraph, - bm); - } else { - for (auto cubeop : bm.getOpsByBlockId(bm.getBlockIdByOp(op))) { - if (!visited[cubeop]) { - visited[cubeop] = true; - passAndCollectCandidates(cubeop, indegree, candidates, visited, - memGraph, bm); - } - } + for (auto [op, _] : indegree) { + if (!bm.isWholeCubeReady(op, indegree) || isFusableOp(op) || visited[op]) { + continue; + } + for (auto cubeop : bm.getOpsInSameBlock(op)) { + if (!visited[cubeop]) { + visited[cubeop] = true; + passAndCollectCandidates(cubeop, indegree, candidates, visited, + memGraph, bm); } } } } -void updateCandidates(Operation *nextFused, - SmallVector &candidates, - DenseMap &indegree, - DenseMap &visited, - const CVPipeline::MemoryDependenceGraph &memGraph) { +static void +updateCandidates(Operation *nextFused, SmallVector &candidates, + DenseMap &indegree, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph) { // 1. Already fuse with nextFused, so remove it from candidates for (auto it = candidates.begin(); it != candidates.end(); it++) { if (*it == nextFused) { @@ -167,34 +141,27 @@ void updateCandidates(Operation *nextFused, } // 2. Add new candidates whose indegree becomes 0 after fusing nextFused. - auto block = nextFused->getBlock(); - SmallVector allusers; - allusers.append(nextFused->getUsers().begin(), nextFused->getUsers().end()); - for (auto memUser : memGraph.getExecAfter(nextFused)) { - allusers.push_back(memUser); - } - for (auto user : allusers) { - auto userInBlock = CVPipeline::getAncestorInBlock(user, block); - if (!userInBlock) { - continue; - } - if (!visited[userInBlock]) { - indegree[userInBlock]--; - if (indegree[userInBlock] == 0 && isFusableOp(userInBlock)) { - visited[userInBlock] = true; - candidates.push_back(userInBlock); + DependencyHelper depHelper{memGraph}; + depHelper.forEachUserInSameBlock(nextFused, [&](Operation *user) { + if (!visited[user]) { + indegree[user]--; + if (indegree[user] == 0 && isFusableOp(user)) { + visited[user] = true; + candidates.push_back(user); } } - } + }); } -void findCandidates(DenseMap &indegree, - SmallVector &candidates, - DenseMap &visited, - const CVPipeline::MemoryDependenceGraph &memGraph, - ComputeBlockIdManager &bm) { +static void findCandidates(DenseMap &indegree, + SmallVector &candidates, + DenseMap &visited, + const CVPipeline::MemoryDependenceGraph &memGraph, + ComputeBlockIdManager &bm) { // 1. if no candidate, try to bypass non-fusable + LOG_DEBUG("Finding source ops............\n"); if (candidates.empty()) { + LOG_DEBUG("No candidates available, try bypass\n"); byPassNonFusable(indegree, candidates, visited, memGraph, bm); } // 2. find candidates whose indegree is 0 and not visited, add them to @@ -205,6 +172,7 @@ void findCandidates(DenseMap &indegree, candidates.push_back(op); } } + LOG_DEBUG("end finding source ops............\n"); } static SmallVector @@ -212,25 +180,17 @@ findOpsAdjacentToCube(Block *block, const SmallVector &fuseGroup, DenseMap &visited, const CVPipeline::MemoryDependenceGraph &memGraph) { SmallVector toProcess; + DependencyHelper depHelper{memGraph}; std::optional blockId; for (Operation *op : fuseGroup) { - SmallVector allUsers; - allUsers.append(op->getUsers().begin(), op->getUsers().end()); - for (auto memUser : memGraph.getExecAfter(op)) { - allUsers.push_back(memUser); - } - - for (auto user : allUsers) { - auto userInBlock = CVPipeline::getAncestorInBlock(user, block); - if (!userInBlock || (block->mightHaveTerminator() && - userInBlock == block->getTerminator())) { - continue; + depHelper.forEachUserInSameBlock(op, [&](Operation *user) { + if (block->mightHaveTerminator() && user == block->getTerminator()) { + return; } - if (!isFusableOp(userInBlock) && !visited[userInBlock]) { + if (!isFusableOp(user) && !visited[user]) { auto newBlockId = getOpBlockId(user); if (!newBlockId.has_value()) { - newBlockId = - getOpBlockId(userInBlock); // Some op will be tagged outside + newBlockId = getOpBlockId(user); // Some op will be tagged outside } if (!blockId.has_value()) { @@ -240,29 +200,17 @@ findOpsAdjacentToCube(Block *block, const SmallVector &fuseGroup, toProcess.push_back(op); } } - } + }); } return toProcess; } -static int getLoopCarriedArgIndex(Value operand, Block *block) { - auto barg = dyn_cast(operand); - if (!barg || barg.getOwner() != block || - !isa(block->getParentOp())) { - return -1; - } - unsigned argIdx = barg.getArgNumber(); - if (argIdx == 0) { - return -1; - } - return argIdx; -} - static SetVector collectKeepOps(Block *block, SmallVector toProcess, const SmallVector &fuseGroup, const CVPipeline::MemoryDependenceGraph &memGraph) { SetVector keepOps; + DependencyHelper depHelper{memGraph}; while (!toProcess.empty()) { Operation *op = toProcess.front(); toProcess.erase(toProcess.begin()); @@ -271,46 +219,22 @@ collectKeepOps(Block *block, SmallVector toProcess, } keepOps.insert(op); - // Add all operands to process - for (auto operand : op->getOperands()) { - if (auto defOp = operand.getDefiningOp()) { - if (!keepOps.contains(defOp) && llvm::is_contained(fuseGroup, defOp)) { - toProcess.push_back(defOp); - } - continue; - } - - // Loop-carried dependency: block argument -> yielded value - int argIdx = getLoopCarriedArgIndex(operand, block); - if (argIdx == -1) { - continue; - } - auto barg = cast(operand); - auto *yieldOp = barg.getOwner()->getTerminator(); - auto *yieldedDef = yieldOp->getOperand(argIdx - 1).getDefiningOp(); - if (!keepOps.contains(yieldedDef) && - llvm::is_contained(fuseGroup, yieldedDef)) { - toProcess.push_back(yieldedDef); - } - } - - // Memory dependency - for (auto memDef : memGraph.getExecBefore(op)) { - if (!keepOps.contains(memDef) && llvm::is_contained(fuseGroup, memDef)) { - toProcess.push_back(memDef); + depHelper.forEachSource(op, [&](Operation *source) { + if (!keepOps.contains(source) && llvm::is_contained(fuseGroup, source)) { + toProcess.push_back(source); } - } + }); } // special case: annotation.mark always follows the defining op - for (auto op : fuseGroup) { + for (auto *op : fuseGroup) { auto markOp = llvm::dyn_cast(op); if (!markOp) { continue; } auto src = markOp.getSrc(); - auto definingOp = src.getDefiningOp(); + auto *definingOp = src.getDefiningOp(); if (definingOp && keepOps.contains(definingOp)) { keepOps.insert(markOp); } @@ -382,21 +306,23 @@ extractToProcessFromFuseGroup(Block *block, } SetVector toRemove; - auto forOp = dyn_cast(block->getParentOp()); - if (forOp) { - for (auto op : nowFuseGroup) { - for (auto operand : op->getOperands()) { - int argIdx = getLoopCarriedArgIndex(operand, block); - if (argIdx <= 0) { - continue; - } - auto *yieldOp = block->getTerminator(); - auto yieldOperand = yieldOp->getOperand(argIdx - 1); - auto *defOp = yieldOperand.getDefiningOp(); - if (defOp && bm.getBlockIdByOp(defOp) == -1 && - !llvm::is_contained(nowFuseGroup, defOp)) { - collectAllUsersInFuseGroup(op, nowFuseGroup, toRemove); - } + auto *terminator = block->getTerminator(); + if (llvm::isa_and_present(terminator) && + isa(block->getParentOp())) { + for (auto *op : nowFuseGroup) { + auto walkResult = op->walk( + [block, terminator, &bm, &nowFuseGroup](Operation *nestedOp) { + for (auto operand : nestedOp->getOperands()) { + auto *defOp = getLoopCarriedDefOp(operand, block); + if (defOp && bm.getBlockIdByOp(defOp) == -1 && + !llvm::is_contained(nowFuseGroup, defOp)) { + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + if (walkResult.wasInterrupted()) { + collectAllUsersInFuseGroup(op, nowFuseGroup, toRemove); } } } @@ -454,19 +380,11 @@ static void evictAndRestoreState( } fuseGroup.assign(keepOps.begin(), keepOps.end()); + DependencyHelper depHelper{memGraph}; // 2. Restore indegree for successors and reset visited for removed ops for (Operation *op : toRemove) { - SmallVector allUsers; - allUsers.append(op->getUsers().begin(), op->getUsers().end()); - for (auto memUser : memGraph.getExecAfter(op)) { - allUsers.push_back(memUser); - } - - for (auto user : allUsers) { - if (auto userInBlock = CVPipeline::getAncestorInBlock(user, block)) { - indegree[userInBlock]++; - } - } + depHelper.forEachUserInSameBlock( + op, [&](Operation *user) { indegree[user]++; }); visited[op] = false; } @@ -489,12 +407,14 @@ static void evictAndRestoreState( } } -void refineFuseGroup(Block *block, SmallVector &nowFuseGroup, - DenseMap &visited, - SmallVector &candidates, - DenseMap &indegree, - const CVPipeline::MemoryDependenceGraph &memGraph, - ComputeBlockIdManager &bm, bool isUBRefineOptEnabled) { +static void refineFuseGroup(Block *block, + SmallVector &nowFuseGroup, + DenseMap &visited, + SmallVector &candidates, + DenseMap &indegree, + const CVPipeline::MemoryDependenceGraph &memGraph, + ComputeBlockIdManager &bm, + bool isUBRefineOptEnabled) { // 1.Find ops in fuse group whose next node is a non-fusable (CUBE-only) op auto toProcess = findOpsAdjacentToCube(block, nowFuseGroup, visited, memGraph); @@ -525,7 +445,7 @@ void refineFuseGroup(Block *block, SmallVector &nowFuseGroup, } // Main function to plan vector block id for one block -llvm::LogicalResult +static llvm::LogicalResult planVectorBlockId(Block *block, const CVPipeline::MemoryDependenceGraph &memGraph, ComputeBlockIdManager &bm, bool isUBRefineOptEnabled) { @@ -533,7 +453,7 @@ planVectorBlockId(Block *block, llvm::DenseMap indegree; llvm::SmallVector queue; llvm::DenseMap visited; // has been visited in search - initializeIndegreeForBlock(block, indegree, memGraph, bm); + initializeIndegreeForBlock(block, indegree, DependencyHelper{memGraph}, bm); // 2. initialize visited and find initial candidates block->walk([&](Operation *op) { @@ -557,10 +477,6 @@ planVectorBlockId(Block *block, updateCandidates(nextFused, queue, indegree, visited, memGraph); } if (queue.empty() || nextFused == nullptr) { - LOG_DEBUG("Prepare to check this group: \n"); - for (auto op : nowFuseGroup) { - LOG_DEBUG("fuseing: " << *op << "\n"); - } // finish one group, assign block id and start next iteration // Cut error operations before assigning block id refineFuseGroup(block, nowFuseGroup, visited, queue, indegree, memGraph, @@ -580,8 +496,22 @@ planVectorBlockId(Block *block, return llvm::success(); } +namespace { + +class PlanVectorBlockPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PlanVectorBlockPass) + + PlanVectorBlockPass() = default; + void runOnOperation() override; + + [[nodiscard]] llvm::StringRef getArgument() const final { + return "plan-vector-block"; + } +}; + void PlanVectorBlockPass::runOnOperation() { - LOG_DEBUG("\n---PlanVectorBlockPass start---\n"); // 1. Build memory dependence graph auto moduleOp = getOperation(); @@ -599,22 +529,35 @@ void PlanVectorBlockPass::runOnOperation() { } // 2. search blocks in topo order and assign block id for each block - auto result = moduleOp.walk([&](Block *block) -> WalkResult { + auto result = moduleOp.walk([&](Block *block) { + if (bm.shouldInheritFromParent(block, CoreType::VECTOR_ONLY)) { + if (llvm::failed(bm.inheritFromParent(block))) { + block->getParentOp()->emitError() + << "[" << DEBUG_TYPE + << "] Sub-blocks failed to inherit block id from parent op"; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + } + if (llvm::failed( planVectorBlockId(block, memDepGraph, bm, isUBRefineOptEnabled))) { return WalkResult::interrupt(); } + return WalkResult::advance(); }); if (result.wasInterrupted()) { - LOG_DEBUG("Failed to plan vector block id for block\n"); - CVPipeline::setFallbackAttr(moduleOp, CVPipeline::ERRCODE_FAILED); + signalPassFailure(); } } +} // namespace + +namespace mlir::triton { + std::unique_ptr> createPlanVectorBlockPass() { return std::make_unique(); } -} // namespace triton -} // namespace mlir +} // namespace mlir::triton diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp index e9a8ff7d85..24c61a72de 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.cpp @@ -33,6 +33,7 @@ #include "llvm/Support/raw_ostream.h" #include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" @@ -42,18 +43,16 @@ #include "mlir/Pass/Pass.h" #include "ascend/include/DynamicCVPipeline/Common/MemoryEffectsTracker.h" +#include "ascend/include/DynamicCVPipeline/Common/Utils.h" #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h" +#include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" #include "ascend/include/DynamicCVPipeline/PlanComputeBlock/ReorderOpsByBlockId.h" -#include "DynamicCVPipeline/Common/Utils.h" -#include "DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h" -#include "TritonToUnstructure/OffsetAnalysis.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - using namespace mlir; static constexpr const char *DEBUG_TYPE = "ReorderOpsByBlockIdPass"; -#define LOG_DEBUG(...) \ - LLVM_DEBUG(llvm::dbgs() << " [" << DEBUG_TYPE << "] " << __VA_ARGS__) + +#define DBGS(...) LLVM_DEBUG(llvm::dbgs() << __VA_ARGS__) +#define LOG_DEBUG(...) DBGS("\n[" << DEBUG_TYPE << "] " << __VA_ARGS__) using namespace triton; using namespace CVPipeline; @@ -117,7 +116,7 @@ void EdgeHelper::addEdge(Operation *pred, Operation *succ) { } if (seen.insert({pred, succ}).second) { LOG_DEBUG("Adding " << (IsMemory ? "memory " : "") << "edge from " << *pred - << " to " << *succ << "\n"); + << " to " << *succ); graph.succs[pred].push_back(succ); graph.preds[succ].push_back(pred); } @@ -135,7 +134,7 @@ BlockOpGraph::BlockOpGraph(ArrayRef allOps, Block *block, EdgeHelper edges(*this, block); for (Operation *op : allOps) { - LOG_DEBUG("Processing op: " << *op << "\n"); + LOG_DEBUG("Processing op: " << *op); // Edges from operand defs (including defs nested inside other ops). for (Value const operand : op->getOperands()) { Operation *defOp = operand.getDefiningOp(); @@ -263,11 +262,11 @@ GroupAdjacencyGraph::GroupAdjacencyGraph( // Logging the constructed group graph. LOG_DEBUG("Group-level edges:\n"); for (unsigned i = 0; i < n; ++i) { - LOG_DEBUG(" Group " << groupIds[i] << " -> "); + DBGS(" Group " << groupIds[i] << " -> "); for (unsigned succIdx : succs[i]) { - LOG_DEBUG(groupIds[succIdx] << " "); + DBGS(groupIds[succIdx] << " "); } - LOG_DEBUG("\n"); + DBGS("\n"); } } @@ -301,9 +300,8 @@ GroupAdjacencyGraph::computeTopologicalOrder() { LOG_DEBUG("Group order: "); for (int id : result) { - LOG_DEBUG(id << " "); + DBGS(id << " "); } - LOG_DEBUG("\n"); if (result.size() == n) { return result; @@ -393,7 +391,6 @@ reorderOpsInBlock(Block &block, const MemoryDependenceGraph &memGraph, } void ReorderOpsByBlockIdPass::runOnOperation() { - LOG_DEBUG("\n=== Pass: TuningOpSeq ===\n"); OpBuilder const builder(&getContext()); auto moduleOp = getOperation(); diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp index 0e835386fe..34253330ec 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp @@ -39,19 +39,6 @@ static constexpr const char *DEBUG_TYPE = "refine-args-block-id"; using namespace mlir::triton; -int getLoopCarriedArgIndex(Value operand, Block *block) { - auto barg = dyn_cast(operand); - if (!barg || barg.getOwner() != block || - !isa(block->getParentOp())) { - return -1; - } - unsigned argIdx = barg.getArgNumber(); - if (argIdx == 0) { - return -1; - } - return argIdx; -} - int findFirstUser(BlockArgument iterArg, Block *forBlock, CVPipeline::ComputeBlockIdManager &bm) { llvm::SetVector visited; @@ -90,9 +77,10 @@ bool isDependenceOther(Operation *yieldDefOp, Block *forBlock, int argsId, } } else { // if have block argument from for block. Skip; - if (getLoopCarriedArgIndex(operand, forBlock) != argsId + 1) { + if (CVPipeline::getLoopCarriedArgIndex(operand, forBlock) != argsId + 1) { LOG_DEBUG("Yield def op depends on other arg:" - << getLoopCarriedArgIndex(operand, forBlock) << "\n"); + << CVPipeline::getLoopCarriedArgIndex(operand, forBlock) + << "\n"); return true; } } From 8418583a5cd15d035f5b9d9e438ada067f22a42a Mon Sep 17 00:00:00 2001 From: chenwuyang Date: Wed, 12 Aug 2026 11:47:48 +0800 Subject: [PATCH 2/4] fix reviews --- .../include/DynamicCVPipeline/Common/Utils.h | 21 ----- .../PlanComputeBlock/Common.h | 8 +- .../PlanComputeBlock/ComputeBlockIdManager.h | 13 ++- .../lib/DynamicCVPipeline/Common/Utils.cpp | 8 +- .../PlanComputeBlock/Common.cpp | 16 ++-- .../ComputeBlockIdManager.cpp | 17 ++-- .../PlanComputeBlock/OpClassifier.cpp | 83 ++++++++++--------- .../PlanComputeBlock/PlanCubeBlock.cpp | 24 +++--- .../PlanComputeBlock/PlanVectorBlockPass.cpp | 15 ++-- .../SplitDataflow/RefineArgsBlockId.cpp | 2 +- .../test_plan_compute_block_regional.mlir | 83 +++++++++++++++++++ ...f22_mlir_four_layer_v2c_inner_dep_outer.py | 2 +- ...f28_mlir_five_layer_v2c_inner_dep_outer.py | 2 +- 13 files changed, 194 insertions(+), 100 deletions(-) create mode 100644 third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PlanComputeBlock/test_plan_compute_block_regional.mlir diff --git a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h index b41290cef0..b10e3ce629 100644 --- a/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h +++ b/third_party/ascend/include/DynamicCVPipeline/Common/Utils.h @@ -86,18 +86,6 @@ constexpr int64_t BYTE_SIZE = 8; static constexpr int crossCoreProducerId = 1; static constexpr int crossCoreConsumerId = 0; -class MoveOnly { -protected: - MoveOnly() = default; - ~MoveOnly() = default; - - MoveOnly(const MoveOnly &) = delete; - MoveOnly &operator=(const MoveOnly &) = delete; - - MoveOnly(MoveOnly &&) = default; - MoveOnly &operator=(MoveOnly &&) = default; -}; - enum CoreType { UNDETERMINED = 0, VECTOR_ONLY = 1 << 0, @@ -143,15 +131,6 @@ inline bool isVectorSimpleOpOrCf(Operation *op) { return getCoreTypeOfSimpleOpOrCf(op) == CoreType::VECTOR_ONLY; } -template -inline llvm::LogicalResult allSucceededShortCircuit(RangeT &&range, - FuncT &&func) { - return llvm::success( - llvm::all_of(std::forward(range), [&](auto &&item) { - return llvm::succeeded(func(std::forward(item))); - })); -} - bool isVectorOnlyOp(Operation *op); bool isScalarLike(Value value); diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h index 4ce5b44ac3..5de5bf5e7f 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/Common.h @@ -57,17 +57,17 @@ class DependencyHelper { void forEachUser(Operation *op, PredFn pred) const; - template + enum class SourceMode { Default, AcrossIterArg }; + template void forEachSource(Operation *op, PredFn pred) const; void forEachUserInSameBlock(Operation *op, PredFn pred) const { forEachUser(op, mapToAncestorInBlock(op->getBlock(), pred)); } - template + template void forEachSourceInSameBlock(Operation *op, PredFn pred) const { - forEachSource(op, - mapToAncestorInBlock(op->getBlock(), pred)); + forEachSource(op, mapToAncestorInBlock(op->getBlock(), pred)); } }; diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h index 4443c740f8..c0a9f06d59 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h @@ -39,7 +39,7 @@ namespace CVPipeline { /** * the class is to promise CUBEID and VECTORID are unified. */ -class ComputeBlockIdManager : MoveOnly { +class ComputeBlockIdManager { public: ComputeBlockIdManager(Operation *root); bool isSameBlock(Operation *a, Operation *b); @@ -54,11 +54,20 @@ class ComputeBlockIdManager : MoveOnly { llvm::LogicalResult inheritFromParent(Block *block); llvm::SmallVector getOpsByBlockId(int blockId) const; + + // Get operations that share the same block_id AND mlir block of op llvm::SmallVector getOpsInSameBlock(Operation *op) const; + std::optional getBlockIdByOpOpt(Operation *op) const; int getNextId(); - int getBlockIdByOp(Operation *op); + int getBlockIdByOp(Operation *op) const; + + ~ComputeBlockIdManager() = default; + ComputeBlockIdManager(const ComputeBlockIdManager &) = delete; + ComputeBlockIdManager &operator=(const ComputeBlockIdManager &) = delete; + ComputeBlockIdManager(ComputeBlockIdManager &&) = delete; + ComputeBlockIdManager &operator=(ComputeBlockIdManager &&) = delete; private: int cntComputeBlockId = 0; diff --git a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp index 283e04fadd..2ab4ddd9f5 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/Common/Utils.cpp @@ -26,9 +26,9 @@ #include "bishengir/Dialect/HIVM/IR/HIVM.h" -static constexpr const char *DEBUG_TYPE = "DynamicCVPipeline/Utils"; +static constexpr const char *DEBUG_TYPE = "dynamic-cv-pipeline-utils"; #define DBGS(...) LLVM_DEBUG(llvm::dbgs() << __VA_ARGS__) -#define LOG_DEBUG(...) DBGS("\n[" << DEBUG_TYPE << "] " << __VA_ARGS__) +#define LOG_DEBUG(...) DBGS("[" << DEBUG_TYPE << "] " << __VA_ARGS__) namespace mlir { namespace CVPipeline { @@ -184,7 +184,9 @@ CoreType getCoreTypeOfSimpleOpOrCf(Operation *op) { return WalkResult::skip(); }); - LOG_DEBUG("CoreType of RegionBranchOp is " << coreType << ": " << *op); + (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; diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp index f13436836a..216d503c40 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/Common.cpp @@ -43,7 +43,11 @@ void DependencyHelper::forEachUser(Operation *op, } } -template +namespace { +using SourceMode = DependencyHelper::SourceMode; +} + +template void DependencyHelper::forEachSource(Operation *op, DependencyHelper::PredFn pred) const { op->walk([&, this, op](Operation *subOp) { @@ -55,7 +59,7 @@ void DependencyHelper::forEachSource(Operation *op, continue; } - if constexpr (AcrossIterArg) { + if constexpr (SM == SourceMode::AcrossIterArg) { if (auto *defOp = getLoopCarriedDefOp(operand, op->getBlock())) { pred(defOp); } @@ -71,11 +75,13 @@ void DependencyHelper::forEachSource(Operation *op, } // Instantiate concrete functions for linking -template void mlir::CVPipeline::DependencyHelper::forEachSource( +template void +mlir::CVPipeline::DependencyHelper::forEachSource( mlir::Operation *op, llvm::function_ref callback) const; -template void mlir::CVPipeline::DependencyHelper::forEachSource( +template void +mlir::CVPipeline::DependencyHelper::forEachSource( mlir::Operation *op, llvm::function_ref callback) const; @@ -85,7 +91,7 @@ void initializeIndegreeForBlock(Block *block, ComputeBlockIdManager &bm) { for (auto *op : llvm::make_pointer_range(block->getOperations())) { indegree[op] = 0; - depHelper.forEachSource(op, [&](Operation *source) { + depHelper.forEachSource(op, [&](Operation *source) { if (source->getBlock() == block && !bm.isSameBlock(source, op)) { indegree[op]++; } diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp index fb5b2d331c..9f95dd325b 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp @@ -43,6 +43,9 @@ ComputeBlockIdManager::ComputeBlockIdManager(Operation *root) { root->walk([&](Operation *op) { if (auto blockIdAttr = op->getAttrOfType(kBlockId)) { auto blockId = blockIdAttr.getInt(); + if (blockId <= 0) { + return; + } opToBlockId[op] = blockId; blockIdToOps[blockId].push_back(op); cntComputeBlockId = @@ -143,7 +146,7 @@ ComputeBlockIdManager::getBlockIdByOpOpt(Operation *op) const { return std::nullopt; } -int ComputeBlockIdManager::getBlockIdByOp(Operation *op) { +int ComputeBlockIdManager::getBlockIdByOp(Operation *op) const { return getBlockIdByOpOpt(op).value_or(-1); } @@ -214,11 +217,13 @@ llvm::LogicalResult ComputeBlockIdManager::inheritFromParent(Block *block) { // 2. if we have marked nested ops with this block id, it could be correct, // since they will be re-marked with the same id, so no failures will be // returned, but this is less robust - return allSucceededShortCircuit( - block->getOperations(), - [this, blockId = blockIdOpt.value()](Operation &op) { - return markAndRecord(&op, blockId); - }); + auto blockId = blockIdOpt.value(); + for (auto &op : *block) { + if (llvm::failed(markAndRecord(&op, blockId))) { + return llvm::failure(); + } + } + return llvm::success(); } } // namespace CVPipeline diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp index b9a3739d08..fc1a089aec 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp @@ -25,6 +25,7 @@ #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Casting.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -36,6 +37,7 @@ #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Interfaces/ViewLikeInterface.h" #include "mlir/Support/LLVM.h" #include "ascend/include/DynamicCVPipeline/Common/Utils.h" @@ -156,6 +158,30 @@ void OpClassifierPass::markCube(Operation *op) { } } +static bool isExtractedLoadStoreRelated(Operation *op) { + if (!op) + return false; + return llvm::TypeSwitch(op) + .Case([](bufferization::ToTensorOp toTensorOp) { + return isExtractedLoadStoreRelated( + toTensorOp.getBuffer().getDefiningOp()); + }) + .Case([](memref::AllocOp allocOp) { + Value memref = allocOp.getMemref(); + for (auto user : memref.getUsers()) { + auto forOp = user->getParentOfType(); + if (forOp && forOp->hasAttr(hivm::ExtractLoadStoreAttr)) + return true; + } + return false; + }) + .Case([](ViewLikeOpInterface viewOp) { + return isExtractedLoadStoreRelated( + viewOp.getViewSource().getDefiningOp()); + }) + .Default([](auto) { return false; }); +} + // ============================================================================ // Pattern: to_tensor → matmul (Upstream) // ============================================================================ @@ -182,24 +208,21 @@ void OpClassifierPass::matchToTensorPattern(Operation *def) { if (!toTensorOp) return; - // special case: implicit transpose -> vector + // special case: implicit transpose -> remains vector if (utils::getAnnotateOpWithAttr(toTensorOp.getResult(), kMayImplicitTransposeWithLastAxis)) { return; } - Value memref = toTensorOp.getBuffer(); - // special case: ExtractLoadStore -> vector - if (llvm::any_of(memref.getUsers(), [](Operation *user) { - auto forOp = user->getParentOfType(); - return forOp && forOp->hasAttr(hivm::ExtractLoadStoreAttr); - })) { + // special case: ExtractedLoadOrStore -> remains vector + if (isExtractedLoadStoreRelated(toTensorOp)) { return; } markCube(toTensorOp); cubeSeeds.push_back(toTensorOp); + Value memref = toTensorOp.getBuffer(); // Also mark the memref allocation as CUBE if (Operation *memrefDef = memref.getDefiningOp()) { markCube(memrefDef); @@ -246,7 +269,7 @@ void OpClassifierPass::matchTransposePattern(Operation *def) { // Helper lambda to check if an operand's defining op qualifies for CUBE seed auto shouldMarkCubeSeed = [](Operation *opDef) -> bool { - if (!opDef) + if (!opDef || isExtractedLoadStoreRelated(opDef)) return false; return (isa(opDef->getDialect()) && !isa(opDef)) || @@ -256,22 +279,16 @@ void OpClassifierPass::matchTransposePattern(Operation *def) { // Check input tensor auto operands = transposeOp->getOperands(); for (const auto &op : operands) { - if (shouldMarkCubeSeed(op.getDefiningOp())) { - markCube(op.getDefiningOp()); - cubeSeeds.push_back(op.getDefiningOp()); - break; // No need to check other operands, one is enough to seed the - // transpose as CUBE + auto defOp = op.getDefiningOp(); + if (!shouldMarkCubeSeed(defOp)) { + continue; } - } - - // Check outs (DpsInits) - auto outs = transposeOp.getDpsInits(); - for (const auto &out : outs) { - if (shouldMarkCubeSeed(out.getDefiningOp())) { - markCube(out.getDefiningOp()); - cubeSeeds.push_back(out.getDefiningOp()); - break; + if (llvm::isa(defOp)) { + matchToTensorPattern(defOp); + continue; } + markCube(defOp); + cubeSeeds.push_back(defOp); } } @@ -756,6 +773,10 @@ int OpClassifierPass::propagateCubeUpstream() { } } + // Skip ExtractedLoadOrStore related op + if (isExtractedLoadStoreRelated(def)) + continue; + // Skip operations inside linalg block (internal values) // But don't skip the linalg op itself if (isInsideNestedLinalgRegion(def)) { @@ -803,24 +824,6 @@ int OpClassifierPass::markRemainingAsVector() { if (opCoreTypes[op] == OP_UNDETERMINED && !isa(op)) { opCoreTypes[op] = OP_VECTOR_ONLY; } - - // ExtractLoadStoreAttr -> force on vector - if (isa(op) && op->hasAttr(hivm::ExtractLoadStoreAttr)) { - op->walk([this](Operation *nestedOp) { - opCoreTypes[nestedOp] = OP_VECTOR_ONLY; - for (auto operand : nestedOp->getOperands()) { - if (auto allocOp = llvm::dyn_cast_if_present( - operand.getDefiningOp())) { - opCoreTypes[allocOp] = OP_VECTOR_ONLY; - for (auto *user : allocOp->getUsers()) { - if (llvm::isa(user)) { - opCoreTypes[user] = OP_VECTOR_ONLY; - } - } - } - } - }); - } } return 0; diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp index ad55d42e75..79a1676aa5 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanCubeBlock.cpp @@ -126,10 +126,11 @@ bool DependencyCycleDetector::detectCycleFrom(Operation *cur) { bool createsCycle = false; depHelper.forEachUserInSameBlock(cur, [&](Operation *user) { - createsCycle = createsCycle || llvm::any_of(bm.getOpsInSameBlock(user), - [this](Operation *user) { - return detectCycleFrom(user); - }); + if (createsCycle) + return; + createsCycle = + llvm::any_of(bm.getOpsInSameBlock(user), + [this](Operation *user) { return detectCycleFrom(user); }); return; }); @@ -184,7 +185,7 @@ void SeedRegionPlanner::run() { size_t head = 0; while (head < group.size()) { Operation *currOp = group[head++]; - depHelper.forEachSource( + depHelper.forEachSource( currOp, [this](Operation *source) { tryAddToGroup(source); }); } } @@ -281,11 +282,14 @@ llvm::LogicalResult TopologicalPartitionPlanner::removeReadyNonCubeOps() { } } if (indegreeBefore == indegree && beforeVisitedSize == bypassVisited.size()) { - if (Operation *parentOp = block->getParentOp()) { - parentOp->emitError("PlanCubeBlock cannot make progress while scheduling " - "cube operations"); - } - dumpQueueAndIndegreeInfo(); + LLVM_DEBUG({ + if (Operation *parentOp = block->getParentOp()) { + LOG_DEBUG("PlanCubeBlock cannot make progress while scheduling " + "cube operations in: " + << *parentOp); + } + dumpQueueAndIndegreeInfo(); + }); return llvm::failure(); } return llvm::success(); diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp index a711d64da5..dbc0d0c496 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/PlanVectorBlockPass.cpp @@ -219,11 +219,13 @@ collectKeepOps(Block *block, SmallVector toProcess, } keepOps.insert(op); - depHelper.forEachSource(op, [&](Operation *source) { - if (!keepOps.contains(source) && llvm::is_contained(fuseGroup, source)) { - toProcess.push_back(source); - } - }); + depHelper.forEachSource( + op, [&](Operation *source) { + if (!keepOps.contains(source) && + llvm::is_contained(fuseGroup, source)) { + toProcess.push_back(source); + } + }); } // special case: annotation.mark always follows the defining op @@ -548,7 +550,8 @@ void PlanVectorBlockPass::runOnOperation() { return WalkResult::advance(); }); if (result.wasInterrupted()) { - signalPassFailure(); + LOG_DEBUG("Failed to plan vector block id for block\n"); + CVPipeline::setFallbackAttr(moduleOp, CVPipeline::ERRCODE_FAILED); } } diff --git a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp index 34253330ec..1b8d961799 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/SplitDataflow/RefineArgsBlockId.cpp @@ -77,7 +77,7 @@ bool isDependenceOther(Operation *yieldDefOp, Block *forBlock, int argsId, } } else { // if have block argument from for block. Skip; - if (CVPipeline::getLoopCarriedArgIndex(operand, forBlock) != argsId + 1) { + if (CVPipeline::getLoopCarriedArgIndex(operand, forBlock) != argsId) { LOG_DEBUG("Yield def op depends on other arg:" << CVPipeline::getLoopCarriedArgIndex(operand, forBlock) << "\n"); diff --git a/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PlanComputeBlock/test_plan_compute_block_regional.mlir b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PlanComputeBlock/test_plan_compute_block_regional.mlir new file mode 100644 index 0000000000..3d148434f8 --- /dev/null +++ b/third_party/ascend/unittest/Conversion/General/DynamicCVPipeline/PlanComputeBlock/test_plan_compute_block_regional.mlir @@ -0,0 +1,83 @@ +// RUN: triton-opt --plan-compute-block %s | FileCheck %s + +module { + // ============================================================================ + // 1. extracted_load_store_stays_vector + // ============================================================================ + // + // CHECK-LABEL: func.func @extracted_load_store_stays_vector( + // CHECK: %[[ALLOC:[A-Za-z0-9_]+]] = memref.alloc() {ssbuffer.block_id = [[B_ID:[0-9]+]] : i32, ssbuffer.core_type = "VECTOR"} + // CHECK: scf.for + // CHECK: %[[SUBVIEW:[A-Za-z0-9_]+]] = memref.subview %[[ALLOC]] + // CHECK: memref.copy %{{.*}}, %[[SUBVIEW]] {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "VECTOR"} + // CHECK: %[[LHS:[A-Za-z0-9_]+]] = bufferization.to_tensor %[[ALLOC]] restrict writable {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "VECTOR"} + // CHECK: linalg.matmul {ssbuffer.block_id = {{[0-9]+}} : i32, ssbuffer.core_type = "CUBE"} + func.func @extracted_load_store_stays_vector( + %arg0: memref, + %rhs: tensor<64x64xf16>) -> tensor<32x64xf32> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %alloc = memref.alloc() : memref<32x64xf16> + + scf.for %arg1 = %c0 to %c32 step %c1 { + %subview = memref.subview %alloc[%arg1, 0] [1, 64] [1, 1] : memref<32x64xf16> to memref<1x64xf16, strided<[64, 1], offset: ?>> + %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [0], sizes: [1, 64], strides: [64, 1] : memref to memref<1x64xf16, strided<[64, 1], offset: ?>> + memref.copy %reinterpret_cast, %subview : memref<1x64xf16, strided<[64, 1], offset: ?>> to memref<1x64xf16, strided<[64, 1], offset: ?>> + } {ExtractedLoadOrStore} + + %lhs = bufferization.to_tensor %alloc restrict writable : memref<32x64xf16> to tensor<32x64xf16> + %out = tensor.empty() : tensor<32x64xf32> + %cst_f32 = arith.constant 0.0 : f32 + %init = linalg.fill ins(%cst_f32 : f32) outs(%out : tensor<32x64xf32>) -> tensor<32x64xf32> + %mm = linalg.matmul ins(%lhs, %rhs : tensor<32x64xf16>, tensor<64x64xf16>) outs(%init : tensor<32x64xf32>) -> tensor<32x64xf32> + return %mm : tensor<32x64xf32> + } + + // ============================================================================ + // 2. cube_control_flow_inheritance + // ============================================================================ + // + // CHECK-LABEL: func.func @cube_control_flow_inheritance( + // CHECK: %[[ALLOC:[A-Za-z0-9_]+]] = memref.alloc() {ssbuffer.block_id = [[B_ID:[0-9]+]] : i32, ssbuffer.core_type = "CUBE"} + // CHECK: scf.if + // CHECK: linalg.fill {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "CUBE"} + // CHECK: %[[LHS:[A-Za-z0-9_]+]] = bufferization.to_tensor %[[ALLOC]] restrict writable {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "CUBE"} + // CHECK: linalg.matmul {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "CUBE"} + func.func @cube_control_flow_inheritance( + %cond: i1, + %arg0: tensor<64x64xf16>, + %arg1: tensor<64x64xf32>) -> tensor<64x64xf32> { + %alloc = memref.alloc() {ssbuffer.core_type = "CUBE"} : memref<64x64xf16> + %cst = arith.constant 0.0 : f16 + + scf.if %cond { + linalg.fill {ssbuffer.core_type = "CUBE"} ins(%cst : f16) outs(%alloc : memref<64x64xf16>) + } + + %lhs = bufferization.to_tensor %alloc restrict writable {ssbuffer.core_type = "CUBE"} : memref<64x64xf16> to tensor<64x64xf16> + %mm = linalg.matmul {ssbuffer.core_type = "CUBE"} ins(%lhs, %arg0 : tensor<64x64xf16>, tensor<64x64xf16>) outs(%arg1 : tensor<64x64xf32>) -> tensor<64x64xf32> + return %mm : tensor<64x64xf32> + } + + // ============================================================================ + // 3. vector_control_flow_inheritance + // ============================================================================ + // + // CHECK-LABEL: func.func @vector_control_flow_inheritance( + // CHECK: %[[ALLOC:[A-Za-z0-9_]+]] = memref.alloc() {ssbuffer.block_id = [[B_ID:[0-9]+]] : i32, ssbuffer.core_type = "VECTOR"} + // CHECK: scf.if + // CHECK: linalg.fill {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "VECTOR"} + // CHECK: memref.copy %[[ALLOC]], %{{.*}} {ssbuffer.block_id = [[B_ID]] : i32, ssbuffer.core_type = "VECTOR"} + func.func @vector_control_flow_inheritance(%cond: i1, %arg0: memref<64x64xf32>) { + %cst = arith.constant 0.0 : f32 + %alloc = memref.alloc() {ssbuffer.core_type = "VECTOR"} : memref<64x64xf32> + + scf.if %cond { + linalg.fill {ssbuffer.core_type = "VECTOR"} ins(%cst : f32) outs(%alloc : memref<64x64xf32>) + } + + memref.copy %alloc, %arg0 {ssbuffer.core_type = "VECTOR"} : memref<64x64xf32> to memref<64x64xf32> + return + } +} diff --git a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf22_mlir_four_layer_v2c_inner_dep_outer.py b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf22_mlir_four_layer_v2c_inner_dep_outer.py index f8dfbc219e..3510417053 100644 --- a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf22_mlir_four_layer_v2c_inner_dep_outer.py +++ b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf22_mlir_four_layer_v2c_inner_dep_outer.py @@ -214,7 +214,7 @@ def test_sdf22(): assert mlir and len(mlir) > 0, "MLIR code generation failed or is empty" assert "func.func @sdf22(" in mlir, \ "Kernel function definition not found in MLIR code" - assert "scope" not in mlir, "Fallback scenario: MLIR code unexpectedly contains the 'scope' keyword" + assert "scope" in mlir, "MLIR code does not contain the 'scope' keyword" # Output MLIR code to the specified path diff --git a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf28_mlir_five_layer_v2c_inner_dep_outer.py b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf28_mlir_five_layer_v2c_inner_dep_outer.py index 7106c0e607..8b3c4df8b3 100644 --- a/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf28_mlir_five_layer_v2c_inner_dep_outer.py +++ b/third_party/ascend/unittest/DynamicCVPipeline_ut/test_sdf28_mlir_five_layer_v2c_inner_dep_outer.py @@ -219,7 +219,7 @@ def test_sdf28(): assert mlir and len(mlir) > 0, "MLIR code generation failed or is empty" assert "func.func @sdf28(" in mlir, \ "Kernel function definition not found in MLIR code" - assert "scope" not in mlir, "Fallback scenario: MLIR code unexpectedly contains the 'scope' keyword" + assert "scope" in mlir, "MLIR code does not contain the 'scope' keyword" # Output MLIR code to the specified path From 719cc8ac4fed7b10dcf90789cfe8bf51a969aecb Mon Sep 17 00:00:00 2001 From: olivervnc Date: Wed, 12 Aug 2026 23:31:56 +0800 Subject: [PATCH 3/4] fix pipeline error: getOpsByBlockId not found when linking --- .../DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h | 2 +- .../PlanComputeBlock/ComputeBlockIdManager.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h index c0a9f06d59..5bc6d593a9 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h @@ -53,7 +53,7 @@ class ComputeBlockIdManager { bool shouldInheritFromParent(Block *block, CoreType requiredCoreType) const; llvm::LogicalResult inheritFromParent(Block *block); - llvm::SmallVector getOpsByBlockId(int blockId) const; + llvm::SmallVector getOpsByBlockId(int blockId); // Get operations that share the same block_id AND mlir block of op llvm::SmallVector getOpsInSameBlock(Operation *op) const; diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp index 9f95dd325b..59c42747ee 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp @@ -107,7 +107,7 @@ void ComputeBlockIdManager::updateBlockId(Operation *op, int blockId) { } llvm::SmallVector -ComputeBlockIdManager::getOpsByBlockId(int blockId) const { +ComputeBlockIdManager::getOpsByBlockId(int blockId) { if (blockId == -1) { return {}; } From e131ad72493364804e817f10c25ff0c81af38981 Mon Sep 17 00:00:00 2001 From: olivervnc Date: Wed, 12 Aug 2026 23:43:15 +0800 Subject: [PATCH 4/4] fix: get block id by op --- .../DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h | 2 +- .../PlanComputeBlock/ComputeBlockIdManager.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h index 5bc6d593a9..091719172b 100644 --- a/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h +++ b/third_party/ascend/include/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.h @@ -61,7 +61,7 @@ class ComputeBlockIdManager { std::optional getBlockIdByOpOpt(Operation *op) const; int getNextId(); - int getBlockIdByOp(Operation *op) const; + int getBlockIdByOp(Operation *op); ~ComputeBlockIdManager() = default; ComputeBlockIdManager(const ComputeBlockIdManager &) = delete; diff --git a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp index 59c42747ee..8eb94f45bc 100644 --- a/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp +++ b/third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/ComputeBlockIdManager.cpp @@ -146,7 +146,7 @@ ComputeBlockIdManager::getBlockIdByOpOpt(Operation *op) const { return std::nullopt; } -int ComputeBlockIdManager::getBlockIdByOp(Operation *op) const { +int ComputeBlockIdManager::getBlockIdByOp(Operation *op) { return getBlockIdByOpOpt(op).value_or(-1); }