From e7e0d22a85935aec87181f45b3a043f6a06abc15 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 11:09:54 -0400 Subject: [PATCH 01/33] Tmp --- test/Proton/scope_id.mlir | 110 ++++++-- .../include/Analysis/ScopeIdAllocation.h | 4 + .../lib/Analysis/ScopeIdAllocation.cpp | 236 ++++++++++++++++-- 3 files changed, 303 insertions(+), 47 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index d49c49882ba6..ff630aa3b830 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -35,21 +35,6 @@ module { proton.record end "name1" tt.return } - - // expected-remark @below {{control_flow}} - tt.func @control_flow(%cond: i1) { - // expected-remark @below {{scope id = 5}} - proton.record start "name0" - scf.if %cond { - // expected-remark @below {{scope id = 6}} - proton.record start "name1" - // expected-remark @below {{scope id = 6}} - proton.record end "name1" - } - // expected-remark @below {{scope id = 5}} - proton.record end "name0" - tt.return - } } // ----- @@ -95,18 +80,103 @@ module { // ----- module { - // expected-remark @below {{condition}} - tt.func @condition(%cond: i1) { + // expected-remark @below {{cf_branch}} + tt.func @cf_branch(%cond: i1) { + ^bb0(%arg0: i1): // expected-remark @below {{scope id = 0}} proton.record start "name0" + cf.cond_br %arg0, ^bb1, ^bb2 + ^bb1: // pred: ^bb0 // expected-remark @below {{scope id = 0}} proton.record end "name0" - scf.if %cond { + cf.br ^bb3 + ^bb2: // pred: ^bb0 + // expected-remark @below {{scope id = 0}} + proton.record end "name0" + cf.br ^bb3 + ^bb3: // preds: ^bb1, ^bb2 + tt.return + } +} + +// ----- + +module { + // expected-remark @below {{cf_reordered}} + tt.func @cf_reordered() { + ^entry: + cf.br ^start + ^exit: + // expected-remark @below {{scope id = 0}} + proton.record end "name0" + tt.return + ^start: + // expected-remark @below {{scope id = 0}} + proton.record start "name0" + cf.br ^exit + } +} + +// ----- + +module { + tt.func @cf_mismatch(%cond: i1) { + ^bb0(%arg0: i1): + proton.record start "name0" + cf.cond_br %arg0, ^bb1, ^bb2 + ^bb1: // pred: ^bb0 + proton.record end "name0" + cf.br ^bb3 + ^bb2: // pred: ^bb0 + cf.br ^bb3 + ^bb3: // preds: ^bb1, ^bb2 + // expected-error @below {{inconsistent proton scope stack across predecessors, expected [name0] but found []}} + tt.return + } +} + +// ----- + +module { + // expected-remark @below {{warp_specialize_balanced}} + tt.func @warp_specialize_balanced() { + // expected-remark @below {{scope id = 0}} + proton.record start "outer" + ttg.warp_specialize() + default { // expected-remark @below {{scope id = 1}} - proton.record start "name0" + proton.record start "default" // expected-remark @below {{scope id = 1}} - proton.record end "name0" + proton.record end "default" + ttg.warp_yield + } + partition0() num_warps(1) { + // expected-remark @below {{scope id = 2}} + proton.record start "partition" + // expected-remark @below {{scope id = 2}} + proton.record end "partition" + ttg.warp_return + } : () -> () + // expected-remark @below {{scope id = 0}} + proton.record end "outer" + tt.return + } +} + +// ----- + +module { + tt.func @warp_specialize_mismatch() { + proton.record start "outer" + ttg.warp_specialize() + default { + ttg.warp_yield } + partition0() num_warps(1) { + proton.record end "outer" + ttg.warp_return + } : () -> () + // expected-error @below {{inconsistent proton scope stack across predecessors, expected [outer] but found []}} tt.return } } diff --git a/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h b/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h index 054c4bdfa45b..718016bd170b 100644 --- a/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h +++ b/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h @@ -14,6 +14,7 @@ namespace mlir { namespace triton::proton { + class ScopeIdAllocation { public: using ScopeId = size_t; @@ -45,7 +46,10 @@ class ScopeIdAllocation { size_t getNumScopes() const { return idToNameMap.size(); } private: + using VirtualBlock = std::pair; + void run(); + void visitTerminator(Operation *op, SmallVector &successors); Operation *funcOp; llvm::DenseMap idToNameMap; diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index c952499684e4..dea13dcc4cff 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -1,3 +1,4 @@ +#include "mlir/Analysis/TopologicalSortUtils.h" #include "Analysis/ScopeIdAllocation.h" #include @@ -9,49 +10,230 @@ namespace triton::proton { #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") +using VirtualBlock = std::pair; + +struct BlockInfo { + llvm::DenseSet activeScopes; + + BlockInfo() = default; + + /// Unions two BlockInfo objects. + void join(const BlockInfo &other) { + for (auto &scope : other.activeScopes) { + this->activeScopes.insert(scope); + } + } + + bool contains(StringRef scopeName) const { + return this->activeScopes.contains(scopeName); + } + + void erase(StringRef scopeName) { + this->activeScopes.erase(scopeName); + } + + void insert(StringRef scopeName) { + this->activeScopes.insert(scopeName); + } + + bool operator ==(const BlockInfo &other) const { + return this->activeScopes == other.activeScopes; + } +}; + void ScopeIdAllocation::run() { - llvm::StringMap nameToIdMap; + DenseMap inputBlockInfoMap; + DenseMap outputBlockInfoMap; + + std::deque blockList; + funcOp->walk([&](Block *block) { + // Start the analysis from the entry blocks of any nested isolated from + // above regions. + if (block->isEntryBlock() && + !isa(block->getParentOp())) + blockList.emplace_back(block); + }); + + // Reachability analysis + while (!blockList.empty()) { + VirtualBlock &virtualBlock = blockList.front(); + blockList.pop_front(); + // Make a copy of the inputblockInfo but not update + auto inputBlockInfo = inputBlockInfoMap[virtualBlock]; + SmallVector successors; + Block::iterator startIt = + virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); + for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { + if (op.hasTrait() || + isa(op)) { + visitTerminator(&op, successors); + break; + } + if (auto recordOp = dyn_cast(&op)) { + auto name = recordOp.getName(); + if (inputBlockInfo.contains(name)) { + if (!recordOp.getIsStart()) { + inputBlockInfo.erase(name); + } + } else { + if (recordOp.getIsStart()) { + inputBlockInfo.insert(name); + } // else don't handle it right now as the scope might be monotonically closed later + } + } + } + // Get the reference because we want to update if it changed + if (outputBlockInfoMap.count(virtualBlock) && + inputBlockInfo == outputBlockInfoMap[virtualBlock]) { + // If we have seen the block before and the inputBlockInfo is the same as + // the outputBlockInfo, we skip the successors + continue; + } + // Update the current block. The block transfer function is not monotonic, + // so overwrite the output state entirely. + outputBlockInfoMap[virtualBlock] = inputBlockInfo; + // Update the successors + for (VirtualBlock &successor : successors) { + inputBlockInfoMap[successor].join(outputBlockInfoMap[virtualBlock]); + blockList.emplace_back(successor); + } + } + + // Go through all blocks, validate reachability analysis results + for (auto iter : inputBlockInfoMap) { + auto &virtualBlock = iter.first; + auto &inputBlockInfo = iter.second; + auto &outputBlockInfo = outputBlockInfoMap[virtualBlock]; + Block::iterator startIt = + virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); + for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { + if (auto recordOp = dyn_cast(&op)) { + auto name = recordOp.getName(); + if (recordOp.getIsStart()) { + if (!outputBlockInfo.contains(name)) { + mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is not closed properly"; + } + } else { + if (!inputBlockInfo.contains(name)) { + mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is closed without being opened"; + } + } + } + } + } + + // Liveness analysis + // For each scope, find its nearest pair + llvm::DenseMap> nameToIdMap; + llvm::DenseMap idToOpMap; std::stack scopeIdStack; - ScopeId id = 0; + ScopeId scopeId = 0; funcOp->walk([&](RecordOp recordOp) { auto name = recordOp.getName(); LDBG("Processing RecordOp: " << recordOp); - if (recordOp.getIsStart()) { - if (!nameToIdMap.contains(name)) { - nameToIdMap[name] = id; - idToNameMap[id] = name; - LDBG("Assigning new scope id " << id << " to name '" << name << "'"); - opToIdMap[recordOp] = id; - if (!scopeIdStack.empty()) { - scopeParentIds.push_back({id, scopeIdStack.top()}); - } - scopeIdStack.push(id); - id++; - } else { - mlir::emitError(recordOp.getLoc(), "The scope name '") - << name << "' must appear in pairs"; - } + if (!nameToIdMap.contains(name)) { + nameToIdMap[name] = {scopeId, /*isStart=*/recordOp.getIsStart()}; + idToNameMap[scopeId] = name; + LDBG("Assigning new scope scopeId " << scopeId << " to op '" << recordOp << "'"); + opToIdMap[recordOp] = scopeId; + idToOpMap[scopeId] = recordOp; + scopeId++; } else { - if (nameToIdMap.contains(name)) { - scopeIdStack.pop(); - opToIdMap[recordOp] = nameToIdMap.lookup(name); - nameToIdMap.erase(name); + auto &[existingId, isStart] = nameToIdMap[name]; + if (isStart == recordOp.getIsStart()) { + // Error: duplicate start or end + mlir::emitError(recordOp.getLoc(), "Scope name '") + << name << "' has duplicate " + << (recordOp.getIsStart() ? "start" : "end") << " record"; } else { - mlir::emitError(recordOp.getLoc(), "The scope name '") - << name << "' must appear in pairs"; + // Matching pair found + LDBG("Found matching pair for scope name '" << name << "' with scopeId " << existingId); + opToIdMap[recordOp] = existingId; + idToOpMap[existingId] = recordOp; + nameToIdMap.erase(name); + scopeIdStack.pop(); } } }); - if (nameToIdMap.size() > 0) { - for (auto &[name, _] : nameToIdMap) { - mlir::emitError(funcOp->getLoc(), "Scope name '") - << name << "' must appear in pairs"; + // Sort all start scopes in the topological order and get the nearest parent + llvm::SetVector startRecordOps; + funcOp->walk([&](RecordOp recordOp) { + if (recordOp.getIsStart()) { + startRecordOps.insert(recordOp); + } + }); + + auto sortedStartRecordOps = mlir::topologicalSort(startRecordOps); + mlir::DominanceInfo domInfo(funcOp); + for (auto i = 0; i < sortedStartRecordOps.size(); ++i) { + auto *op = sortedStartRecordOps[i]; + for (auto j = 0; j < i; ++j) { + auto *maybeParentOp = sortedStartRecordOps[j]; + if (domInfo.dominates(maybeParentOp, op)) { + auto parentId = opToIdMap.lookup(maybeParentOp); + auto childId = opToIdMap.lookup(op); + scopeParentIds.push_back({childId, parentId}); + break; + } } } } +void ScopeIdAllocation::visitTerminator( + Operation *op, SmallVector &successors) { + if (isa(op)) { + // Collect the block successors of the branch. + for (Block *successor : op->getSuccessors()) + successors.emplace_back(successor, Block::iterator()); + return; + } + + if (auto br = dyn_cast(op)) { + // The successors of an operation with regions can be queried via an + // interface. The operation branches to the entry blocks of its region + // successors. It can also branch to after itself. + SmallVector regions; + br.getSuccessorRegions(RegionBranchPoint::parent(), regions); + for (RegionSuccessor ®ion : regions) { + if (region.isParent()) { + successors.emplace_back(br->getBlock(), br->getIterator()); + } else { + Block &block = region.getSuccessor()->front(); + successors.emplace_back(&block, Block::iterator()); + } + } + return; + } + + // FIXME: `ReturnLike` adds `RegionBranchTerminatorOpInterface` for some + // reason. Check that the parent is actually a `RegionBranchOpInterface`. + auto br = dyn_cast(op); + if (br && isa(br->getParentOp())) { + // Check the successors of a region branch terminator. It can branch to + // another region of its parent operation or to after the parent op. + SmallVector operands(br->getNumOperands()); + SmallVector regions; + br.getSuccessorRegions(operands, regions); + for (RegionSuccessor ®ion : regions) { + if (region.isParent()) { + Operation *parent = br->getParentOp(); + successors.emplace_back(parent->getBlock(), parent->getIterator()); + } else { + Block &block = region.getSuccessor()->front(); + successors.emplace_back(&block, Block::iterator()); + } + } + return; + } + + // Otherwise, it could be a return op + if (op->hasTrait()) + return; + llvm_unreachable("Unknown terminator encountered in membar analysis"); +} + ModuleScopeIdAllocation::ModuleScopeIdAllocation(ModuleOp moduleOp) : CallGraph(moduleOp) { ScopeIdAllocation::ScopeId funcScopeId = 0; From fa816becb1a61ce3a8714d991c47bc02d4cb430d Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 15:18:10 -0400 Subject: [PATCH 02/33] Fix --- .../include/Analysis/ScopeIdAllocation.h | 3 + .../lib/Analysis/ScopeIdAllocation.cpp | 63 +++++++++++++++---- 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h b/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h index 718016bd170b..11bd9e57f9f7 100644 --- a/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h +++ b/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h @@ -49,6 +49,9 @@ class ScopeIdAllocation { using VirtualBlock = std::pair; void run(); + void reachability(); + void liveness(); + void dominance(); void visitTerminator(Operation *op, SmallVector &successors); Operation *funcOp; diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index dea13dcc4cff..00bdde8c6e1d 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -42,23 +42,57 @@ struct BlockInfo { }; void ScopeIdAllocation::run() { + // Stage the analysis to match downstream consumers of scope metadata: + // + // - reachability(): Track active scopes at CFG boundaries and flag malformed + // lifetimes. Example MLIR: + // scf.if %cond { + // proton.record start @"foo" + // } + // Because `"foo"` never ends on the `then` branch, reachability() emits + // `The scope name 'foo' is not closed properly`. + // + // - liveness(): Pair start/end records and assign a shared numeric ID. Example + // MLIR: + // proton.record start @"foo" + // … + // proton.record end @"foo" + // Both ops are mapped to the same ScopeId in `opToIdMap`. + // + // - dominance(): Infer the parent/child hierarchy between scopes via + // dominance. Example MLIR: + // proton.record start @"outer" + // scf.if %cond { + // proton.record start @"inner" + // … + // proton.record end @"inner" + // } + // proton.record end @"outer" + // Because the start of `"outer"` dominates `"inner"`, dominance() records + // `(innerId -> outerId)` in `scopeParentIds`. + reachability(); + liveness(); + dominance(); +} + +void ScopeIdAllocation::reachability() { DenseMap inputBlockInfoMap; DenseMap outputBlockInfoMap; - std::deque blockList; + std::deque virtualBlockList; funcOp->walk([&](Block *block) { // Start the analysis from the entry blocks of any nested isolated from // above regions. if (block->isEntryBlock() && !isa(block->getParentOp())) - blockList.emplace_back(block); + virtualBlockList.emplace_back(block, Block::iterator()); }); - // Reachability analysis - while (!blockList.empty()) { - VirtualBlock &virtualBlock = blockList.front(); - blockList.pop_front(); - // Make a copy of the inputblockInfo but not update + while (!virtualBlockList.empty()) { + VirtualBlock &virtualBlock = virtualBlockList.front(); + virtualBlockList.pop_front(); + // Evaluate the transfer function for this block starting from the cached + // input state. auto inputBlockInfo = inputBlockInfoMap[virtualBlock]; SmallVector successors; Block::iterator startIt = @@ -95,7 +129,7 @@ void ScopeIdAllocation::run() { // Update the successors for (VirtualBlock &successor : successors) { inputBlockInfoMap[successor].join(outputBlockInfoMap[virtualBlock]); - blockList.emplace_back(successor); + virtualBlockList.emplace_back(successor); } } @@ -122,8 +156,11 @@ void ScopeIdAllocation::run() { } } - // Liveness analysis - // For each scope, find its nearest pair +} + +void ScopeIdAllocation::liveness() { + // Stage 2: pair start/end records that refer to the same scope name and + // assign a numeric ID that downstream passes can reuse. llvm::DenseMap> nameToIdMap; llvm::DenseMap idToOpMap; std::stack scopeIdStack; @@ -157,7 +194,11 @@ void ScopeIdAllocation::run() { } }); - // Sort all start scopes in the topological order and get the nearest parent +} + +void ScopeIdAllocation::dominance() { + // Stage 3: determine parentage between scopes by checking dominance of start + // operations. llvm::SetVector startRecordOps; funcOp->walk([&](RecordOp recordOp) { if (recordOp.getIsStart()) { From d88c59f691ee8f892f7097595c32d1d16123b2a7 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 15:24:06 -0400 Subject: [PATCH 03/33] Fix --- .../proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 00bdde8c6e1d..1bbe7454df77 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -123,9 +123,8 @@ void ScopeIdAllocation::reachability() { // the outputBlockInfo, we skip the successors continue; } - // Update the current block. The block transfer function is not monotonic, - // so overwrite the output state entirely. - outputBlockInfoMap[virtualBlock] = inputBlockInfo; + // Update the current block + outputBlockInfoMap[virtualBlock].join(inputBlockInfo); // Update the successors for (VirtualBlock &successor : successors) { inputBlockInfoMap[successor].join(outputBlockInfoMap[virtualBlock]); @@ -137,7 +136,7 @@ void ScopeIdAllocation::reachability() { for (auto iter : inputBlockInfoMap) { auto &virtualBlock = iter.first; auto &inputBlockInfo = iter.second; - auto &outputBlockInfo = outputBlockInfoMap[virtualBlock]; + auto &outputBlockInfo = outputBlockInfoMap[virtualBlock]; Block::iterator startIt = virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { From 9f62119a1eb9b4c986d4d88ece80979994ffa6b8 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 15:55:48 -0400 Subject: [PATCH 04/33] Update --- .../lib/Analysis/ScopeIdAllocation.cpp | 33 ++++++++++++++----- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 1bbe7454df77..e74958ee3ca5 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -39,6 +39,14 @@ struct BlockInfo { bool operator ==(const BlockInfo &other) const { return this->activeScopes == other.activeScopes; } + + void dump() const { + auto &err = llvm::errs(); + err << "Active Scopes:\n"; + for (auto &scope : activeScopes) { + err << " " << scope << "\n"; + } + } }; void ScopeIdAllocation::run() { @@ -135,26 +143,35 @@ void ScopeIdAllocation::reachability() { // Go through all blocks, validate reachability analysis results for (auto iter : inputBlockInfoMap) { auto &virtualBlock = iter.first; - auto &inputBlockInfo = iter.second; - auto &outputBlockInfo = outputBlockInfoMap[virtualBlock]; + auto inputBlockInfo = iter.second; + auto outputBlockInfo = outputBlockInfoMap[virtualBlock]; + DenseSet unclosedScopes; Block::iterator startIt = virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { if (auto recordOp = dyn_cast(&op)) { auto name = recordOp.getName(); if (recordOp.getIsStart()) { - if (!outputBlockInfo.contains(name)) { - mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is not closed properly"; - } + inputBlockInfo.insert(name); + unclosedScopes.insert(name); } else { - if (!inputBlockInfo.contains(name)) { - mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is closed without being opened"; + if (inputBlockInfo.contains(name)) { + inputBlockInfo.erase(name); + unclosedScopes.erase(name); + } else { + mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is ended without being opened"; } } } } + for (auto &scopeName : unclosedScopes) { + if (!outputBlockInfo.contains(scopeName)) { + mlir::emitError(virtualBlock.first->getParentOp()->getLoc(), + "The scope name '") + << scopeName << "' is not closed properly"; + } + } } - } void ScopeIdAllocation::liveness() { From a71774085a832fb2cd93bda25fbab0fd626dafb4 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 16:10:11 -0400 Subject: [PATCH 05/33] Update --- test/Proton/scope_id.mlir | 37 ++++++++++--------------------------- 1 file changed, 10 insertions(+), 27 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index ff630aa3b830..5c1245c62ea2 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -82,15 +82,14 @@ module { module { // expected-remark @below {{cf_branch}} tt.func @cf_branch(%cond: i1) { - ^bb0(%arg0: i1): // expected-remark @below {{scope id = 0}} proton.record start "name0" - cf.cond_br %arg0, ^bb1, ^bb2 - ^bb1: // pred: ^bb0 + cf.cond_br %cond, ^bb1, ^bb2 + ^bb1: // pred: ^entry // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^bb3 - ^bb2: // pred: ^bb0 + ^bb2: // pred: ^entry // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^bb3 @@ -120,17 +119,19 @@ module { // ----- module { + // expected-remark @below {{cf_mismatch}} tt.func @cf_mismatch(%cond: i1) { - ^bb0(%arg0: i1): + // expected-remark @below {{scope id = 0}} proton.record start "name0" - cf.cond_br %arg0, ^bb1, ^bb2 - ^bb1: // pred: ^bb0 + cf.cond_br %cond, ^bb1, ^bb2 + ^bb1: // pred: ^entry + // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^bb3 - ^bb2: // pred: ^bb0 + ^bb2: // pred: ^entry cf.br ^bb3 ^bb3: // preds: ^bb1, ^bb2 - // expected-error @below {{inconsistent proton scope stack across predecessors, expected [name0] but found []}} + // expected-error @below {{inconsistent proton scope stack across predecessors, expected [] but found [name0]}} tt.return } } @@ -162,21 +163,3 @@ module { tt.return } } - -// ----- - -module { - tt.func @warp_specialize_mismatch() { - proton.record start "outer" - ttg.warp_specialize() - default { - ttg.warp_yield - } - partition0() num_warps(1) { - proton.record end "outer" - ttg.warp_return - } : () -> () - // expected-error @below {{inconsistent proton scope stack across predecessors, expected [outer] but found []}} - tt.return - } -} From c90b70ed5360f4a326aa50d622c2509238c6950a Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 16:12:47 -0400 Subject: [PATCH 06/33] Update --- .../proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index e74958ee3ca5..06c4a5057cad 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -210,6 +210,15 @@ void ScopeIdAllocation::liveness() { } }); + if (!nameToIdMap.empty()) { + for (auto &[name, idIsStartPair] : nameToIdMap) { + auto &[id, isStart] = idIsStartPair; + auto unclosedOp = idToOpMap.lookup(id); + mlir::emitError(unclosedOp.getLoc(), "Scope name '") + << name << "' is not properly closed (missing " + << (isStart ? "end" : "start") << " record)"; + } + } } void ScopeIdAllocation::dominance() { From 46c8ff0df7ab8f692403a4daa30b3269c10b5932 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 16:20:16 -0400 Subject: [PATCH 07/33] Update --- test/Proton/scope_id.mlir | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 5c1245c62ea2..51cea50afa91 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -80,17 +80,14 @@ module { // ----- module { - // expected-remark @below {{cf_branch}} tt.func @cf_branch(%cond: i1) { - // expected-remark @below {{scope id = 0}} proton.record start "name0" cf.cond_br %cond, ^bb1, ^bb2 ^bb1: // pred: ^entry - // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^bb3 ^bb2: // pred: ^entry - // expected-remark @below {{scope id = 0}} + // expected-error@+1 {{scope 'name0' was not properly closed (missing start record)}} proton.record end "name0" cf.br ^bb3 ^bb3: // preds: ^bb1, ^bb2 From dc3e31ca7d81944254900beefb6e3703f925af68 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 16:21:17 -0400 Subject: [PATCH 08/33] Fix --- test/Proton/scope_id.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 51cea50afa91..647f4b46f25c 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -87,7 +87,7 @@ module { proton.record end "name0" cf.br ^bb3 ^bb2: // pred: ^entry - // expected-error@+1 {{scope 'name0' was not properly closed (missing start record)}} + // expected-error@+1 {{Scope name 'name0' was not properly closed (missing start record)}} proton.record end "name0" cf.br ^bb3 ^bb3: // preds: ^bb1, ^bb2 From 94f3c784848dd31e9fd0f3ea72f27f43a2b1d52b Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 16:22:43 -0400 Subject: [PATCH 09/33] Fix --- test/Proton/scope_id.mlir | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 647f4b46f25c..440b272624a7 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -80,14 +80,17 @@ module { // ----- module { + // expected-remark @below {{cf_branch}} tt.func @cf_branch(%cond: i1) { + // expected-remark @below {{scope id = 0}} proton.record start "name0" cf.cond_br %cond, ^bb1, ^bb2 ^bb1: // pred: ^entry + // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^bb3 ^bb2: // pred: ^entry - // expected-error@+1 {{Scope name 'name0' was not properly closed (missing start record)}} + // expected-error@+1 {{Scope name 'name0' is not properly closed (missing start record)}} proton.record end "name0" cf.br ^bb3 ^bb3: // preds: ^bb1, ^bb2 From 7a530af6fd01e7c7c4741b865d7c5e88bf482dc9 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 16:23:34 -0400 Subject: [PATCH 10/33] Fix --- test/Proton/scope_id.mlir | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 440b272624a7..881536a142b5 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -80,17 +80,14 @@ module { // ----- module { - // expected-remark @below {{cf_branch}} + // expected-error@+1 {{Scope name 'name0' is not properly closed (missing start record)}} tt.func @cf_branch(%cond: i1) { - // expected-remark @below {{scope id = 0}} proton.record start "name0" cf.cond_br %cond, ^bb1, ^bb2 ^bb1: // pred: ^entry - // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^bb3 ^bb2: // pred: ^entry - // expected-error@+1 {{Scope name 'name0' is not properly closed (missing start record)}} proton.record end "name0" cf.br ^bb3 ^bb3: // preds: ^bb1, ^bb2 From e27c08b7ee3c9d781a45ee645d11e10f8ac08390 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 21:18:46 -0400 Subject: [PATCH 11/33] Update --- test/Proton/scope_id.mlir | 38 ------------------------------- test/Proton/scope_id_invalid.mlir | 35 ++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 38 deletions(-) create mode 100644 test/Proton/scope_id_invalid.mlir diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 881536a142b5..8a50b9fcf0e0 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -79,24 +79,6 @@ module { // ----- -module { - // expected-error@+1 {{Scope name 'name0' is not properly closed (missing start record)}} - tt.func @cf_branch(%cond: i1) { - proton.record start "name0" - cf.cond_br %cond, ^bb1, ^bb2 - ^bb1: // pred: ^entry - proton.record end "name0" - cf.br ^bb3 - ^bb2: // pred: ^entry - proton.record end "name0" - cf.br ^bb3 - ^bb3: // preds: ^bb1, ^bb2 - tt.return - } -} - -// ----- - module { // expected-remark @below {{cf_reordered}} tt.func @cf_reordered() { @@ -115,26 +97,6 @@ module { // ----- -module { - // expected-remark @below {{cf_mismatch}} - tt.func @cf_mismatch(%cond: i1) { - // expected-remark @below {{scope id = 0}} - proton.record start "name0" - cf.cond_br %cond, ^bb1, ^bb2 - ^bb1: // pred: ^entry - // expected-remark @below {{scope id = 0}} - proton.record end "name0" - cf.br ^bb3 - ^bb2: // pred: ^entry - cf.br ^bb3 - ^bb3: // preds: ^bb1, ^bb2 - // expected-error @below {{inconsistent proton scope stack across predecessors, expected [] but found [name0]}} - tt.return - } -} - -// ----- - module { // expected-remark @below {{warp_specialize_balanced}} tt.func @warp_specialize_balanced() { diff --git a/test/Proton/scope_id_invalid.mlir b/test/Proton/scope_id_invalid.mlir new file mode 100644 index 000000000000..85aa32c02f98 --- /dev/null +++ b/test/Proton/scope_id_invalid.mlir @@ -0,0 +1,35 @@ +// RUN: triton-opt --split-input-file --test-print-scope-id-allocation -verify-diagnostics=error -o /dev/null %s + +module { + // expected-error@below {{Scope name 'name0' is not properly closed (missing start record)}} + tt.func @cf_branch(%cond: i1) { + proton.record start "name0" + cf.cond_br %cond, ^bb1, ^bb2 + ^bb1: // pred: ^entry + proton.record end "name0" + cf.br ^bb3 + ^bb2: // pred: ^entry + proton.record end "name0" + cf.br ^bb3 + ^bb3: // preds: ^bb1, ^bb2 + tt.return + } +} + +// ----- + +module { + tt.func @cf_mismatch(%cond: i1) { + ^entry: + proton.record start "name0" + cf.cond_br %cond, ^then, ^else + ^then: // pred: ^entry + proton.record end "name0" + cf.br ^merge + ^else: // pred: ^entry + cf.br ^merge + ^merge: // preds: ^then, ^else + // expected-error @below {{inconsistent proton scope stack across predecessors, expected [] but found [name0]}} + tt.return + } +} From d131c79d5bf6a85c8681eb2a2497e7057a32f8ac Mon Sep 17 00:00:00 2001 From: Jokeren Date: Thu, 30 Oct 2025 21:36:09 -0400 Subject: [PATCH 12/33] Take out --- test/Proton/scope_id_invalid.mlir | 35 ------------------------------- 1 file changed, 35 deletions(-) delete mode 100644 test/Proton/scope_id_invalid.mlir diff --git a/test/Proton/scope_id_invalid.mlir b/test/Proton/scope_id_invalid.mlir deleted file mode 100644 index 85aa32c02f98..000000000000 --- a/test/Proton/scope_id_invalid.mlir +++ /dev/null @@ -1,35 +0,0 @@ -// RUN: triton-opt --split-input-file --test-print-scope-id-allocation -verify-diagnostics=error -o /dev/null %s - -module { - // expected-error@below {{Scope name 'name0' is not properly closed (missing start record)}} - tt.func @cf_branch(%cond: i1) { - proton.record start "name0" - cf.cond_br %cond, ^bb1, ^bb2 - ^bb1: // pred: ^entry - proton.record end "name0" - cf.br ^bb3 - ^bb2: // pred: ^entry - proton.record end "name0" - cf.br ^bb3 - ^bb3: // preds: ^bb1, ^bb2 - tt.return - } -} - -// ----- - -module { - tt.func @cf_mismatch(%cond: i1) { - ^entry: - proton.record start "name0" - cf.cond_br %cond, ^then, ^else - ^then: // pred: ^entry - proton.record end "name0" - cf.br ^merge - ^else: // pred: ^entry - cf.br ^merge - ^merge: // preds: ^then, ^else - // expected-error @below {{inconsistent proton scope stack across predecessors, expected [] but found [name0]}} - tt.return - } -} From 5c58f5806e15ae09a663e65b5788c96358c5446d Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 09:44:39 -0400 Subject: [PATCH 13/33] Update --- test/Proton/scope_id.mlir | 99 ++++++++++++++++++- .../lib/Analysis/ScopeIdAllocation.cpp | 4 + 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 8a50b9fcf0e0..28e07bf717da 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -1,4 +1,4 @@ -// RUN: triton-opt --split-input-file --test-print-scope-id-allocation -verify-diagnostics -o /dev/null %s +// RUN: triton-opt --split-input-file --test-print-scope-id-allocation -verify-diagnostics=only-expected -o /dev/null %s module { // expected-remark @below {{one_scope}} @@ -95,6 +95,26 @@ module { } } +// ----- + +module { + // expected-remark @below {{cf_single_branch}} + tt.func @cf_single_branch(%cond: i1) { + // expected-remark @below {{scope id = 0}} + proton.record start "name0" + cf.cond_br %cond, ^then, ^else + ^then: // pred: ^entry + // expected-remark @below {{scope id = 0}} + proton.record end "name0" + cf.br ^merge + ^else: // pred: ^entry + cf.br ^merge + ^merge: // preds: ^then, ^else + tt.return + } +} + + // ----- module { @@ -122,3 +142,80 @@ module { tt.return } } + +// ----- + +module { + // expected-remark @below {{cf_liveness_error}} + tt.func @cf_liveness_error(%cond: i1) { + // expected-remark @below {{scope id = 0}} + proton.record start "name0" + cf.cond_br %cond, ^then, ^else + ^then: // pred: ^entry + // expected-remark @below {{scope id = 0}} + proton.record end "name0" + cf.br ^merge + ^else: // pred: ^entry + // expected-remark @below {{scope id = 0}} + proton.record end "name0" + cf.br ^merge + ^merge: // preds: ^then, ^else + tt.return + } +} + +// ----- + +module { + tt.func @cf_unclosed() { + proton.record start "unclosed" + } +} + +// ----- + +module { + tt.func @cf_dangling_end() { + // expected-error @below {{The scope name 'dangling' is ended without being opened}} + proton.record end "dangling" + tt.return + } +} + +// ----- + +module { + tt.func @cf_branch_unclosed_dangling(%cond: i1) { + cf.cond_br %cond, ^then, ^else + ^then: // pred: ^entry + // expected-error @below {{The scope name 'ghost_then' is ended without being opened}} + proton.record start "ghost" + cf.br ^merge + ^else: // pred: ^entry + // expected-error @below {{The scope name 'ghost_else' is ended without being opened}} + proton.record end "ghost" + cf.br ^merge + ^merge: // preds: ^then, ^else + tt.return + } +} + +// ----- + +module { + tt.func @cf_merge_unclosed(%cond: i1) { + cf.cond_br %cond, ^then, ^else + proton.record start "ghost" + ^then: // pred: ^entry + // expected-error @below {{The scope name 'ghost_then' is ended without being opened}} + proton.record stop "ghost" + cf.br ^merge + ^else: // pred: ^entry + // expected-error @below {{The scope name 'ghost_else' is ended without being opened}} + proton.record start "ghost" + cf.br ^merge + ^merge: // preds: ^then, ^else + proton.record end "ghost" + tt.return + } +} \ No newline at end of file diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 06c4a5057cad..03a01be8dda1 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -152,6 +152,10 @@ void ScopeIdAllocation::reachability() { if (auto recordOp = dyn_cast(&op)) { auto name = recordOp.getName(); if (recordOp.getIsStart()) { + if (inputBlockInfo.contains(name)) { + mlir::emitError(recordOp.getLoc(), "The scope name '") + << name << "' is started without being ended"; + } inputBlockInfo.insert(name); unclosedScopes.insert(name); } else { From a6ab3b5a4234864719b5729bd8b13b6927af23f5 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 13:09:22 -0400 Subject: [PATCH 14/33] Fix --- test/Proton/scope_id.mlir | 95 +++++++++++++++++-- .../lib/Analysis/ScopeIdAllocation.cpp | 36 ++++--- 2 files changed, 114 insertions(+), 17 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 28e07bf717da..5a823dbace36 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -145,6 +145,48 @@ module { // ----- +module { + tt.func @cf_loop_closed() { + ^entry: + %c0 = arith.constant 0 : index + cf.br ^loop(%c0 : index) + ^exit: + tt.return + ^loop(%iv: index): + proton.record start "loop_body" + %c1 = arith.constant 1 : index + %next = arith.addi %iv, %c1 + %c2 = arith.constant 2 : index + %cond = arith.cmpi ult, %next, %c2 + proton.record end "loop_body" + cf.cond_br %cond, ^loop(%next : index), ^exit + } +} + +// ----- + +module { + tt.func @cf_loop_closed_two_blocks() { + ^entry: + %c0 = arith.constant 0 : index + cf.br ^loop(%c0 : index) + ^exit: + tt.return + ^loop(%iv: index): + proton.record start "loop_body" + %c1 = arith.constant 1 : index + %next = arith.addi %iv, %c1 + cf.br ^loop_body(%next : index) + ^loop_body(%iv_next: index): + %c2 = arith.constant 2 : index + %cond = arith.cmpi ult, %iv_next, %c2 + proton.record end "loop_body" + cf.cond_br %cond, ^loop(%iv_next : index), ^exit + } +} + +// ----- + module { // expected-remark @below {{cf_liveness_error}} tt.func @cf_liveness_error(%cond: i1) { @@ -167,6 +209,7 @@ module { // ----- module { + // expected-error @below {{The scope name 'unclosed' is started without being closed}} tt.func @cf_unclosed() { proton.record start "unclosed" } @@ -175,8 +218,8 @@ module { // ----- module { + // expected-error @below {{The scope name 'dangling' is closed without being opened}} tt.func @cf_dangling_end() { - // expected-error @below {{The scope name 'dangling' is ended without being opened}} proton.record end "dangling" tt.return } @@ -185,14 +228,14 @@ module { // ----- module { + // expected-error @below {{The scope name 'ghost' is started without being closed}} + // expected-error @below {{The scope name 'ghost' is closed without being opened}} tt.func @cf_branch_unclosed_dangling(%cond: i1) { cf.cond_br %cond, ^then, ^else ^then: // pred: ^entry - // expected-error @below {{The scope name 'ghost_then' is ended without being opened}} proton.record start "ghost" cf.br ^merge ^else: // pred: ^entry - // expected-error @below {{The scope name 'ghost_else' is ended without being opened}} proton.record end "ghost" cf.br ^merge ^merge: // preds: ^then, ^else @@ -203,19 +246,59 @@ module { // ----- module { + // expected-error @below {{The scope name 'ghost' is started without being closed}} tt.func @cf_merge_unclosed(%cond: i1) { cf.cond_br %cond, ^then, ^else proton.record start "ghost" ^then: // pred: ^entry - // expected-error @below {{The scope name 'ghost_then' is ended without being opened}} proton.record stop "ghost" cf.br ^merge ^else: // pred: ^entry - // expected-error @below {{The scope name 'ghost_else' is ended without being opened}} proton.record start "ghost" cf.br ^merge ^merge: // preds: ^then, ^else proton.record end "ghost" tt.return } -} \ No newline at end of file +} + +// ----- + +module { + // expected-error @below {{The scope name 'loop' is started without being closed}} + tt.func @cf_loop_unclosed() { + ^entry: + %c0 = arith.constant 0 : index + cf.br ^loop(%c0 : index) + ^exit: + tt.return + ^loop(%iv: index): + proton.record start "loop" + %c1 = arith.constant 1 : index + %next = arith.addi %iv, %c1 + %c2 = arith.constant 2 : index + %cond = arith.cmpi ult, %next, %c2 + cf.cond_br %cond, ^loop(%next : index), ^exit + } +} + +// ----- + +module { + // expected-error @below {{The scope name 'loop' has end record that dominates its start record}} + tt.func @cf_loop_end_before_start() { + ^entry: + %c0 = arith.constant 0 : index + cf.br ^loop(%c0 : index) + ^exit: + tt.return + ^loop(%iv: index): + proton.record end "loop" + %c1 = arith.constant 1 : index + %next = arith.addi %iv, %c1 + %c2 = arith.constant 2 : index + %cond = arith.cmpi ult, %next, %c2 + proton.record start "loop" + cf.cond_br %cond, ^loop(%next : index), ^exit + } +} diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 03a01be8dda1..71391d3f4436 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -1,8 +1,6 @@ #include "mlir/Analysis/TopologicalSortUtils.h" #include "Analysis/ScopeIdAllocation.h" -#include - namespace mlir { namespace triton::proton { @@ -154,7 +152,7 @@ void ScopeIdAllocation::reachability() { if (recordOp.getIsStart()) { if (inputBlockInfo.contains(name)) { mlir::emitError(recordOp.getLoc(), "The scope name '") - << name << "' is started without being ended"; + << name << "' is started without being closed"; } inputBlockInfo.insert(name); unclosedScopes.insert(name); @@ -163,7 +161,7 @@ void ScopeIdAllocation::reachability() { inputBlockInfo.erase(name); unclosedScopes.erase(name); } else { - mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is ended without being opened"; + mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is closed without being opened"; } } } @@ -183,7 +181,6 @@ void ScopeIdAllocation::liveness() { // assign a numeric ID that downstream passes can reuse. llvm::DenseMap> nameToIdMap; llvm::DenseMap idToOpMap; - std::stack scopeIdStack; ScopeId scopeId = 0; funcOp->walk([&](RecordOp recordOp) { @@ -209,7 +206,6 @@ void ScopeIdAllocation::liveness() { opToIdMap[recordOp] = existingId; idToOpMap[existingId] = recordOp; nameToIdMap.erase(name); - scopeIdStack.pop(); } } }); @@ -228,15 +224,33 @@ void ScopeIdAllocation::liveness() { void ScopeIdAllocation::dominance() { // Stage 3: determine parentage between scopes by checking dominance of start // operations. - llvm::SetVector startRecordOps; + mlir::DominanceInfo domInfo(funcOp); + llvm::DenseMap startRecordMap; + llvm::DenseMap endRecordMap; funcOp->walk([&](RecordOp recordOp) { - if (recordOp.getIsStart()) { - startRecordOps.insert(recordOp); - } + auto scopeId = opToIdMap.lookup(recordOp); + if (recordOp.getIsStart()) + startRecordMap[scopeId] = recordOp.getOperation(); + else + endRecordMap[scopeId] = recordOp.getOperation(); }); + for (auto &[scopeId, startOp] : startRecordMap) { + auto *endOp = endRecordMap.lookup(scopeId); + if (!endOp) + continue; + if (domInfo.dominates(endOp, startOp)) { + auto name = idToNameMap.lookup(scopeId); + mlir::emitError(endOp->getLoc(), "The scope name '") + << name << "' has end record that dominates its start record"; + } + } + + llvm::SetVector startRecordOps; + for (auto &[scopeId, startOp] : startRecordMap) { + startRecordOps.insert(startOp); + } auto sortedStartRecordOps = mlir::topologicalSort(startRecordOps); - mlir::DominanceInfo domInfo(funcOp); for (auto i = 0; i < sortedStartRecordOps.size(); ++i) { auto *op = sortedStartRecordOps[i]; for (auto j = 0; j < i; ++j) { From ca211400bfb6438ad54d6ab2a96e0a928c27c664 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 13:13:36 -0400 Subject: [PATCH 15/33] Fix --- test/Proton/scope_id.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 5a823dbace36..565b64b3bab9 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -188,7 +188,7 @@ module { // ----- module { - // expected-remark @below {{cf_liveness_error}} + // expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}} tt.func @cf_liveness_error(%cond: i1) { // expected-remark @below {{scope id = 0}} proton.record start "name0" From 384904000572a9eec65488d51f7a5b07823336f9 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 13:14:50 -0400 Subject: [PATCH 16/33] Fix --- test/Proton/scope_id.mlir | 3 --- third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp | 4 ++-- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 565b64b3bab9..5cd554e93cca 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -190,15 +190,12 @@ module { module { // expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}} tt.func @cf_liveness_error(%cond: i1) { - // expected-remark @below {{scope id = 0}} proton.record start "name0" cf.cond_br %cond, ^then, ^else ^then: // pred: ^entry - // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^merge ^else: // pred: ^entry - // expected-remark @below {{scope id = 0}} proton.record end "name0" cf.br ^merge ^merge: // preds: ^then, ^else diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 71391d3f4436..a470e10e8431 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -197,7 +197,7 @@ void ScopeIdAllocation::liveness() { auto &[existingId, isStart] = nameToIdMap[name]; if (isStart == recordOp.getIsStart()) { // Error: duplicate start or end - mlir::emitError(recordOp.getLoc(), "Scope name '") + mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' has duplicate " << (recordOp.getIsStart() ? "start" : "end") << " record"; } else { @@ -214,7 +214,7 @@ void ScopeIdAllocation::liveness() { for (auto &[name, idIsStartPair] : nameToIdMap) { auto &[id, isStart] = idIsStartPair; auto unclosedOp = idToOpMap.lookup(id); - mlir::emitError(unclosedOp.getLoc(), "Scope name '") + mlir::emitError(unclosedOp.getLoc(), "The scope name '") << name << "' is not properly closed (missing " << (isStart ? "end" : "start") << " record)"; } From 01062d22c9403886ae6400f602ff3deff576a7a8 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 13:39:06 -0400 Subject: [PATCH 17/33] Fix --- test/Proton/scope_id.mlir | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 5cd554e93cca..6a6868c78778 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -155,7 +155,7 @@ module { ^loop(%iv: index): proton.record start "loop_body" %c1 = arith.constant 1 : index - %next = arith.addi %iv, %c1 + %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index %cond = arith.cmpi ult, %next, %c2 proton.record end "loop_body" @@ -175,7 +175,7 @@ module { ^loop(%iv: index): proton.record start "loop_body" %c1 = arith.constant 1 : index - %next = arith.addi %iv, %c1 + %next = arith.addi %iv, %c1 : index cf.br ^loop_body(%next : index) ^loop_body(%iv_next: index): %c2 = arith.constant 2 : index @@ -206,17 +206,18 @@ module { // ----- module { - // expected-error @below {{The scope name 'unclosed' is started without being closed}} tt.func @cf_unclosed() { + // expected-error @below {{The scope name 'unclosed' is started without being closed}} proton.record start "unclosed" + tt.return } } // ----- module { - // expected-error @below {{The scope name 'dangling' is closed without being opened}} tt.func @cf_dangling_end() { + // expected-error @below {{The scope name 'dangling' is closed without being opened}} proton.record end "dangling" tt.return } @@ -225,14 +226,14 @@ module { // ----- module { - // expected-error @below {{The scope name 'ghost' is started without being closed}} - // expected-error @below {{The scope name 'ghost' is closed without being opened}} tt.func @cf_branch_unclosed_dangling(%cond: i1) { cf.cond_br %cond, ^then, ^else ^then: // pred: ^entry + // expected-error @below {{The scope name 'ghost' is started without being closed}} proton.record start "ghost" cf.br ^merge ^else: // pred: ^entry + // expected-error @below {{The scope name 'ghost' is closed without being opened}} proton.record end "ghost" cf.br ^merge ^merge: // preds: ^then, ^else @@ -243,7 +244,6 @@ module { // ----- module { - // expected-error @below {{The scope name 'ghost' is started without being closed}} tt.func @cf_merge_unclosed(%cond: i1) { cf.cond_br %cond, ^then, ^else proton.record start "ghost" @@ -251,6 +251,7 @@ module { proton.record stop "ghost" cf.br ^merge ^else: // pred: ^entry + // expected-error @below {{The scope name 'ghost' is started without being closed}} proton.record start "ghost" cf.br ^merge ^merge: // preds: ^then, ^else @@ -272,7 +273,7 @@ module { ^loop(%iv: index): proton.record start "loop" %c1 = arith.constant 1 : index - %next = arith.addi %iv, %c1 + %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index %cond = arith.cmpi ult, %next, %c2 cf.cond_br %cond, ^loop(%next : index), ^exit @@ -292,7 +293,7 @@ module { ^loop(%iv: index): proton.record end "loop" %c1 = arith.constant 1 : index - %next = arith.addi %iv, %c1 + %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index %cond = arith.cmpi ult, %next, %c2 proton.record start "loop" From 17645877e8b217f7a41796bbe73f91581c9ad6f6 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 14:01:23 -0400 Subject: [PATCH 18/33] Fix --- test/Proton/scope_id.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 6a6868c78778..fba3e3d92482 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -188,7 +188,6 @@ module { // ----- module { - // expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}} tt.func @cf_liveness_error(%cond: i1) { proton.record start "name0" cf.cond_br %cond, ^then, ^else @@ -196,6 +195,7 @@ module { proton.record end "name0" cf.br ^merge ^else: // pred: ^entry + // expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}} proton.record end "name0" cf.br ^merge ^merge: // preds: ^then, ^else From 89082cfddb56a7b94d76dc39fd30e1ad4e54e580 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 14:03:28 -0400 Subject: [PATCH 19/33] Fix --- test/Proton/scope_id.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index fba3e3d92482..eacd513d3e82 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -207,7 +207,7 @@ module { module { tt.func @cf_unclosed() { - // expected-error @below {{The scope name 'unclosed' is started without being closed}} + // expected-error @below {{The scope name 'unclosed' is not properly closed (missing end record)}} proton.record start "unclosed" tt.return } From 304746c913e8b40e25ab5911ed6894ebf58d86b5 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 14:29:25 -0400 Subject: [PATCH 20/33] Update --- test/Proton/scope_id.mlir | 1 - 1 file changed, 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index eacd513d3e82..182c37eab426 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -229,7 +229,6 @@ module { tt.func @cf_branch_unclosed_dangling(%cond: i1) { cf.cond_br %cond, ^then, ^else ^then: // pred: ^entry - // expected-error @below {{The scope name 'ghost' is started without being closed}} proton.record start "ghost" cf.br ^merge ^else: // pred: ^entry From b4b105853486e3a648acb2294262c9cd31619475 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 14:35:56 -0400 Subject: [PATCH 21/33] Fix --- test/Proton/scope_id.mlir | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 182c37eab426..6476855b6e71 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -244,12 +244,14 @@ module { module { tt.func @cf_merge_unclosed(%cond: i1) { - cf.cond_br %cond, ^then, ^else + cf.br ^start(%cond : i1) + ^start(%cond_arg: i1): proton.record start "ghost" - ^then: // pred: ^entry - proton.record stop "ghost" + cf.cond_br %cond_arg, ^then, ^else + ^then: // pred: ^start + proton.record end "ghost" cf.br ^merge - ^else: // pred: ^entry + ^else: // pred: ^start // expected-error @below {{The scope name 'ghost' is started without being closed}} proton.record start "ghost" cf.br ^merge @@ -264,7 +266,6 @@ module { module { // expected-error @below {{The scope name 'loop' is started without being closed}} tt.func @cf_loop_unclosed() { - ^entry: %c0 = arith.constant 0 : index cf.br ^loop(%c0 : index) ^exit: @@ -284,7 +285,6 @@ module { module { // expected-error @below {{The scope name 'loop' has end record that dominates its start record}} tt.func @cf_loop_end_before_start() { - ^entry: %c0 = arith.constant 0 : index cf.br ^loop(%c0 : index) ^exit: From c6f7486495cc95baf32d60f222c2105db20afc9b Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 14:38:45 -0400 Subject: [PATCH 22/33] Update --- test/Proton/scope_id.mlir | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 6476855b6e71..e25d48ec68dd 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -264,13 +264,13 @@ module { // ----- module { - // expected-error @below {{The scope name 'loop' is started without being closed}} tt.func @cf_loop_unclosed() { %c0 = arith.constant 0 : index cf.br ^loop(%c0 : index) ^exit: tt.return ^loop(%iv: index): + // expected-error @below {{The scope name 'loop' is started without being closed}} proton.record start "loop" %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index @@ -283,13 +283,13 @@ module { // ----- module { - // expected-error @below {{The scope name 'loop' has end record that dominates its start record}} tt.func @cf_loop_end_before_start() { %c0 = arith.constant 0 : index cf.br ^loop(%c0 : index) ^exit: tt.return ^loop(%iv: index): + // expected-error @below {{The scope name 'loop' has end record that dominates its start record}} proton.record end "loop" %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index From 9932cc3619475fa7a0ad3348eace3323f94bfd93 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 14:45:36 -0400 Subject: [PATCH 23/33] Fix --- test/Proton/scope_id.mlir | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index e25d48ec68dd..17a594cc54a5 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -157,7 +157,7 @@ module { %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index - %cond = arith.cmpi ult, %next, %c2 + %cond = arith.cmpi ult, %next, %c2: index proton.record end "loop_body" cf.cond_br %cond, ^loop(%next : index), ^exit } @@ -179,7 +179,7 @@ module { cf.br ^loop_body(%next : index) ^loop_body(%iv_next: index): %c2 = arith.constant 2 : index - %cond = arith.cmpi ult, %iv_next, %c2 + %cond = arith.cmpi ult, %iv_next, %c2: index proton.record end "loop_body" cf.cond_br %cond, ^loop(%iv_next : index), ^exit } @@ -275,7 +275,7 @@ module { %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index - %cond = arith.cmpi ult, %next, %c2 + %cond = arith.cmpi ult, %next, %c2: index cf.cond_br %cond, ^loop(%next : index), ^exit } } @@ -294,7 +294,7 @@ module { %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index - %cond = arith.cmpi ult, %next, %c2 + %cond = arith.cmpi ult, %next, %c2: index proton.record start "loop" cf.cond_br %cond, ^loop(%next : index), ^exit } From a796349d9bd406f4e8122f64ea16bb3ebbfaa571 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 15:06:59 -0400 Subject: [PATCH 24/33] Fix --- test/Proton/scope_id.mlir | 28 +++++++++++++++++++++++ test/lib/Proton/TestScopeIdAllocation.cpp | 14 ++++++++++++ 2 files changed, 42 insertions(+) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 17a594cc54a5..1dba89be674d 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -4,8 +4,10 @@ module { // expected-remark @below {{one_scope}} tt.func @one_scope() { // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" tt.return } @@ -13,12 +15,16 @@ module { // expected-remark @below {{two_scopes}} tt.func @two_scopes() { // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" // expected-remark @below {{scope id = 2}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name1" // expected-remark @below {{scope id = 2}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name1" tt.return } @@ -26,12 +32,16 @@ module { // expected-remark @below {{two_scopes_overlap}} tt.func @two_scopes_overlap() { // expected-remark @below {{scope id = 3}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 4}} + // expected-remark @below {{scope parent id = 3}} proton.record start "name1" // expected-remark @below {{scope id = 3}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" // expected-remark @below {{scope id = 4}} + // expected-remark @below {{scope parent id = 3}} proton.record end "name1" tt.return } @@ -43,8 +53,10 @@ module { // expected-remark @below {{inner}} tt.func @inner() { // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" tt.return } @@ -52,9 +64,11 @@ module { // expected-remark @below {{outer}} tt.func @outer() { // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" tt.call @inner() : () -> () // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" tt.return } @@ -66,12 +80,16 @@ module { // expected-remark @below {{duplicate}} tt.func @duplicate() { // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" tt.return } @@ -86,10 +104,12 @@ module { cf.br ^start ^exit: // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" tt.return ^start: // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" cf.br ^exit } @@ -101,10 +121,12 @@ module { // expected-remark @below {{cf_single_branch}} tt.func @cf_single_branch(%cond: i1) { // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name0" cf.cond_br %cond, ^then, ^else ^then: // pred: ^entry // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "name0" cf.br ^merge ^else: // pred: ^entry @@ -121,23 +143,29 @@ module { // expected-remark @below {{warp_specialize_balanced}} tt.func @warp_specialize_balanced() { // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "outer" ttg.warp_specialize() default { // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = 0}} proton.record start "default" // expected-remark @below {{scope id = 1}} + // expected-remark @below {{scope parent id = 0}} proton.record end "default" ttg.warp_yield } partition0() num_warps(1) { // expected-remark @below {{scope id = 2}} + // expected-remark @below {{scope parent id = 0}} proton.record start "partition" // expected-remark @below {{scope id = 2}} + // expected-remark @below {{scope parent id = 0}} proton.record end "partition" ttg.warp_return } : () -> () // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "outer" tt.return } diff --git a/test/lib/Proton/TestScopeIdAllocation.cpp b/test/lib/Proton/TestScopeIdAllocation.cpp index 7140c0508ea4..8862d18f12c8 100644 --- a/test/lib/Proton/TestScopeIdAllocation.cpp +++ b/test/lib/Proton/TestScopeIdAllocation.cpp @@ -1,6 +1,8 @@ #include "mlir/Pass/Pass.h" #include "third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h" +#include + using namespace mlir; using namespace triton::proton; @@ -30,9 +32,21 @@ struct TestScopeIdAllocationPass moduleOp.walk([&](triton::FuncOp funcOp) { auto opName = SymbolTable::getSymbolName(funcOp).getValue().str(); mlir::emitRemark(funcOp.getLoc(), opName); + llvm::DenseMap + parentScopeIdMap; + for (auto [childId, parentId] : + moduleScopeIdAllocation.getScopeIdParents(funcOp)) { + parentScopeIdMap.insert({childId, parentId}); + } funcOp.walk([&](RecordOp recordOp) { auto scopeId = moduleScopeIdAllocation.getOpScopeId(recordOp); mlir::emitRemark(recordOp.getLoc()) << "scope id = " << scopeId; + int64_t parentId = -1; + if (auto parentIt = parentScopeIdMap.find(scopeId); + parentIt != parentScopeIdMap.end()) + parentId = static_cast(parentIt->second); + mlir::emitRemark(recordOp.getLoc()) + << "scope parent id = " << parentId; }); }); } From 226d17b90bd71768b425dacfb09fe3f683d6dddd Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 15:18:31 -0400 Subject: [PATCH 25/33] Update --- .../proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index a470e10e8431..37d39f5b912a 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -253,9 +253,12 @@ void ScopeIdAllocation::dominance() { auto sortedStartRecordOps = mlir::topologicalSort(startRecordOps); for (auto i = 0; i < sortedStartRecordOps.size(); ++i) { auto *op = sortedStartRecordOps[i]; - for (auto j = 0; j < i; ++j) { + for (int j = i; j >= 0; --j) { auto *maybeParentOp = sortedStartRecordOps[j]; - if (domInfo.dominates(maybeParentOp, op)) { + auto scopeId = opToIdMap.lookup(op); + auto endRecordOp = endRecordMap.lookup(scopeId); + if (domInfo.dominates(maybeParentOp, op) && + domInfo.dominates(op, endRecordOp)) { auto parentId = opToIdMap.lookup(maybeParentOp); auto childId = opToIdMap.lookup(op); scopeParentIds.push_back({childId, parentId}); From 90cd0776ffaae6302f12393493dcb80ddb70de96 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 15:21:36 -0400 Subject: [PATCH 26/33] Fix --- third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 37d39f5b912a..a2205cb6aa83 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -251,9 +251,9 @@ void ScopeIdAllocation::dominance() { startRecordOps.insert(startOp); } auto sortedStartRecordOps = mlir::topologicalSort(startRecordOps); - for (auto i = 0; i < sortedStartRecordOps.size(); ++i) { + for (int i = 0; i < sortedStartRecordOps.size(); ++i) { auto *op = sortedStartRecordOps[i]; - for (int j = i; j >= 0; --j) { + for (int j = i - 1; j >= 0; --j) { auto *maybeParentOp = sortedStartRecordOps[j]; auto scopeId = opToIdMap.lookup(op); auto endRecordOp = endRecordMap.lookup(scopeId); From 225cdcb576688d8f6b77f7665474d8eacbda5edc Mon Sep 17 00:00:00 2001 From: Jokeren Date: Fri, 31 Oct 2025 15:25:59 -0400 Subject: [PATCH 27/33] Fix --- third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index a2205cb6aa83..a049cba08eb2 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -255,7 +255,7 @@ void ScopeIdAllocation::dominance() { auto *op = sortedStartRecordOps[i]; for (int j = i - 1; j >= 0; --j) { auto *maybeParentOp = sortedStartRecordOps[j]; - auto scopeId = opToIdMap.lookup(op); + auto scopeId = opToIdMap.lookup(maybeParentOp); auto endRecordOp = endRecordMap.lookup(scopeId); if (domInfo.dominates(maybeParentOp, op) && domInfo.dominates(op, endRecordOp)) { From 08b337450f6c99e878320129ff9df71f6f1d48e6 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Sat, 1 Nov 2025 09:51:52 -0400 Subject: [PATCH 28/33] Improve --- test/Proton/scope_id.mlir | 65 +++++++++++++------ .../lib/Analysis/ScopeIdAllocation.cpp | 52 ++++++++++----- 2 files changed, 82 insertions(+), 35 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 1dba89be674d..1674f6455b14 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -35,14 +35,31 @@ module { // expected-remark @below {{scope parent id = -1}} proton.record start "name0" // expected-remark @below {{scope id = 4}} - // expected-remark @below {{scope parent id = 3}} + // expected-remark @below {{scope parent id = -1}} proton.record start "name1" // expected-remark @below {{scope id = 3}} // expected-remark @below {{scope parent id = -1}} proton.record end "name0" // expected-remark @below {{scope id = 4}} - // expected-remark @below {{scope parent id = 3}} + // expected-remark @below {{scope parent id = -1}} + proton.record end "name1" + tt.return + } + + // expected-remark @below {{nested_scopes}} + tt.func @nested_scopes() { + // expected-remark @below {{scope id = 5}} + // expected-remark @below {{scope parent id = -1}} + proton.record start "name0" + // expected-remark @below {{scope id = 6}} + // expected-remark @below {{scope parent id = 5}} + proton.record start "name1" + // expected-remark @below {{scope id = 6}} + // expected-remark @below {{scope parent id = 5}} proton.record end "name1" + // expected-remark @below {{scope id = 5}} + // expected-remark @below {{scope parent id = -1}} + proton.record end "name0" tt.return } } @@ -174,6 +191,7 @@ module { // ----- module { + // expected-remark @below {{cf_loop_closed}} tt.func @cf_loop_closed() { ^entry: %c0 = arith.constant 0 : index @@ -181,11 +199,15 @@ module { ^exit: tt.return ^loop(%iv: index): + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "loop_body" %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index %c2 = arith.constant 2 : index %cond = arith.cmpi ult, %next, %c2: index + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "loop_body" cf.cond_br %cond, ^loop(%next : index), ^exit } @@ -194,6 +216,7 @@ module { // ----- module { + // expected-remark @below {{cf_loop_closed_two_blocks}} tt.func @cf_loop_closed_two_blocks() { ^entry: %c0 = arith.constant 0 : index @@ -201,6 +224,8 @@ module { ^exit: tt.return ^loop(%iv: index): + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "loop_body" %c1 = arith.constant 1 : index %next = arith.addi %iv, %c1 : index @@ -208,6 +233,8 @@ module { ^loop_body(%iv_next: index): %c2 = arith.constant 2 : index %cond = arith.cmpi ult, %iv_next, %c2: index + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "loop_body" cf.cond_br %cond, ^loop(%iv_next : index), ^exit } @@ -216,17 +243,9 @@ module { // ----- module { - tt.func @cf_liveness_error(%cond: i1) { - proton.record start "name0" - cf.cond_br %cond, ^then, ^else - ^then: // pred: ^entry - proton.record end "name0" - cf.br ^merge - ^else: // pred: ^entry - // expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}} - proton.record end "name0" - cf.br ^merge - ^merge: // preds: ^then, ^else + tt.func @cf_unclosed() { + // expected-error @below {{The scope name 'unclosed' is not properly closed (missing end record)}} + proton.record start "unclosed" tt.return } } @@ -234,9 +253,9 @@ module { // ----- module { - tt.func @cf_unclosed() { - // expected-error @below {{The scope name 'unclosed' is not properly closed (missing end record)}} - proton.record start "unclosed" + tt.func @cf_dangling_end() { + // expected-error @below {{The scope name 'dangling' is closed without being opened}} + proton.record end "dangling" tt.return } } @@ -244,9 +263,17 @@ module { // ----- module { - tt.func @cf_dangling_end() { - // expected-error @below {{The scope name 'dangling' is closed without being opened}} - proton.record end "dangling" + tt.func @cf_liveness_error(%cond: i1) { + proton.record start "name0" + cf.cond_br %cond, ^then, ^else + ^then: // pred: ^entry + proton.record end "name0" + cf.br ^merge + ^else: // pred: ^entry + // expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}} + proton.record end "name0" + cf.br ^merge + ^merge: // preds: ^then, ^else tt.return } } diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index a049cba08eb2..318d7f85543b 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -51,21 +51,46 @@ void ScopeIdAllocation::run() { // Stage the analysis to match downstream consumers of scope metadata: // // - reachability(): Track active scopes at CFG boundaries and flag malformed - // lifetimes. Example MLIR: + // lifetimes. We optimistically collect all the "potentially" unclosed and + // closing scopes here, and validate them after the dataflow converges. + // + // Starting example MLIR: // scf.if %cond { // proton.record start @"foo" // } // Because `"foo"` never ends on the `then` branch, reachability() emits // `The scope name 'foo' is not closed properly`. // - // - liveness(): Pair start/end records and assign a shared numeric ID. Example + // Valid example MLIR: + // scf.if %cond { + // proton.record start @"foo" + // } + // proton.record end @"foo" + // + // We don't emit errors because we assume scf.if could be executed and it's + // up to the user to ensure proper semantics. + // + // - liveness(): Pair start/end records and assign a shared numeric ID. + // For each start record, we look for the nearest closing scope with the same name for pairing. + // + // Example // MLIR: // proton.record start @"foo" // … // proton.record end @"foo" // Both ops are mapped to the same ScopeId in `opToIdMap`. // - // - dominance(): Infer the parent/child hierarchy between scopes via + // + // - dominance(): + // (1) Check the dominance of start/end records to ensure well-formedness. + // Example MLIR: + // proton.record end @"foo" + // … + // proton.record start @"foo" + // + // Because the end of `"foo"` dominates its start, dominance() emits an error. + // + // (2) Infer the parent/child hierarchy between scopes via // dominance. Example MLIR: // proton.record start @"outer" // scf.if %cond { @@ -166,13 +191,6 @@ void ScopeIdAllocation::reachability() { } } } - for (auto &scopeName : unclosedScopes) { - if (!outputBlockInfo.contains(scopeName)) { - mlir::emitError(virtualBlock.first->getParentOp()->getLoc(), - "The scope name '") - << scopeName << "' is not closed properly"; - } - } } } @@ -253,13 +271,15 @@ void ScopeIdAllocation::dominance() { auto sortedStartRecordOps = mlir::topologicalSort(startRecordOps); for (int i = 0; i < sortedStartRecordOps.size(); ++i) { auto *op = sortedStartRecordOps[i]; + auto scopeId = opToIdMap.lookup(op); + auto endOp = endRecordMap.lookup(scopeId); for (int j = i - 1; j >= 0; --j) { - auto *maybeParentOp = sortedStartRecordOps[j]; - auto scopeId = opToIdMap.lookup(maybeParentOp); - auto endRecordOp = endRecordMap.lookup(scopeId); - if (domInfo.dominates(maybeParentOp, op) && - domInfo.dominates(op, endRecordOp)) { - auto parentId = opToIdMap.lookup(maybeParentOp); + auto *parentStartOp = sortedStartRecordOps[j]; + auto parentScopeId = opToIdMap.lookup(parentStartOp); + auto parentEndOp = endRecordMap.lookup(parentScopeId); + if (domInfo.dominates(parentStartOp, op) && + domInfo.dominates(endOp, parentEndOp)) { + auto parentId = opToIdMap.lookup(parentStartOp); auto childId = opToIdMap.lookup(op); scopeParentIds.push_back({childId, parentId}); break; From a0bc75d734f89929e1210624e4b97cca9602fece Mon Sep 17 00:00:00 2001 From: Jokeren Date: Sat, 1 Nov 2025 16:29:23 -0400 Subject: [PATCH 29/33] Fix --- test/Proton/scope_id.mlir | 48 ++++- .../lib/Analysis/ScopeIdAllocation.cpp | 164 +++++++++--------- 2 files changed, 126 insertions(+), 86 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 1674f6455b14..0e1afce6739c 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -165,19 +165,19 @@ module { ttg.warp_specialize() default { // expected-remark @below {{scope id = 1}} - // expected-remark @below {{scope parent id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "default" // expected-remark @below {{scope id = 1}} - // expected-remark @below {{scope parent id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "default" ttg.warp_yield } partition0() num_warps(1) { // expected-remark @below {{scope id = 2}} - // expected-remark @below {{scope parent id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record start "partition" // expected-remark @below {{scope id = 2}} - // expected-remark @below {{scope parent id = 0}} + // expected-remark @below {{scope parent id = -1}} proton.record end "partition" ttg.warp_return } : () -> () @@ -307,7 +307,6 @@ module { proton.record end "ghost" cf.br ^merge ^else: // pred: ^start - // expected-error @below {{The scope name 'ghost' is started without being closed}} proton.record start "ghost" cf.br ^merge ^merge: // preds: ^then, ^else @@ -354,3 +353,42 @@ module { cf.cond_br %cond, ^loop(%next : index), ^exit } } + +// ----- + +module { + tt.func @cf_if_unclosed(%cond: i1) { + scf.if %cond { + // expected-error @below {{The scope name 'if_only' is not properly closed (missing end record)}} + proton.record start "if_only" + } + tt.return + } +} + +// ----- + +module { + tt.func @cf_duplicate_start() { + // expected-error @below {{The scope name 'dup_scope' is not properly closed (missing end record)}} + proton.record start "dup_scope" + // expected-error @below {{The scope name 'dup_scope' is started without being closed}} + // expected-error @below {{The scope name 'dup_scope' has duplicate start record}} + proton.record start "dup_scope" + tt.return + } +} + +// ----- + +module { + tt.func @cf_duplicate_end() { + // expected-error @below {{The scope name 'dup_scope' is closed without being opened}} + // expected-error @below {{The scope name 'dup_scope' is not properly closed (missing start record)}} + proton.record end "dup_scope" + // expected-error @below {{The scope name 'dup_scope' is closed without being opened}} + // expected-error @below {{The scope name 'dup_scope' has duplicate end record}} + proton.record end "dup_scope" + tt.return + } +} diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 318d7f85543b..738e124acf4a 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -11,7 +11,9 @@ namespace triton::proton { using VirtualBlock = std::pair; struct BlockInfo { - llvm::DenseSet activeScopes; + using ScopeId = ScopeIdAllocation::ScopeId; + + llvm::DenseSet activeScopes; BlockInfo() = default; @@ -22,16 +24,16 @@ struct BlockInfo { } } - bool contains(StringRef scopeName) const { - return this->activeScopes.contains(scopeName); + bool contains(ScopeId scopeId) const { + return this->activeScopes.contains(scopeId); } - void erase(StringRef scopeName) { - this->activeScopes.erase(scopeName); + void erase(ScopeId scopeId) { + this->activeScopes.erase(scopeId); } - void insert(StringRef scopeName) { - this->activeScopes.insert(scopeName); + void insert(ScopeId scopeId) { + this->activeScopes.insert(scopeId); } bool operator ==(const BlockInfo &other) const { @@ -50,6 +52,21 @@ struct BlockInfo { void ScopeIdAllocation::run() { // Stage the analysis to match downstream consumers of scope metadata: // + // - liveness(): Pair start/end records and assign a shared numeric ID. + // There are multiple ways to pair start/end records: + // We choose a simple approach here: + // For each start record, we look for the nearest closing scope with the same name for pairing. + // + // Example + // MLIR: + // proton.record start @"foo" <- scopeId = 0 + // … + // proton.record end @"foo" <- scopeId = 0 + // … + // proton.record start @"foo" <- scopeId = 1 + // … + // proton.record end @"foo" <- scopeId = 1 + // // - reachability(): Track active scopes at CFG boundaries and flag malformed // lifetimes. We optimistically collect all the "potentially" unclosed and // closing scopes here, and validate them after the dataflow converges. @@ -70,17 +87,6 @@ void ScopeIdAllocation::run() { // We don't emit errors because we assume scf.if could be executed and it's // up to the user to ensure proper semantics. // - // - liveness(): Pair start/end records and assign a shared numeric ID. - // For each start record, we look for the nearest closing scope with the same name for pairing. - // - // Example - // MLIR: - // proton.record start @"foo" - // … - // proton.record end @"foo" - // Both ops are mapped to the same ScopeId in `opToIdMap`. - // - // // - dominance(): // (1) Check the dominance of start/end records to ensure well-formedness. // Example MLIR: @@ -101,11 +107,54 @@ void ScopeIdAllocation::run() { // proton.record end @"outer" // Because the start of `"outer"` dominates `"inner"`, dominance() records // `(innerId -> outerId)` in `scopeParentIds`. - reachability(); liveness(); + reachability(); dominance(); } +void ScopeIdAllocation::liveness() { + llvm::DenseMap> nameToIdMap; + llvm::DenseMap idToOpMap; + ScopeId scopeId = 0; + + funcOp->walk([&](RecordOp recordOp) { + auto name = recordOp.getName(); + LDBG("Processing RecordOp: " << recordOp); + if (!nameToIdMap.contains(name)) { + nameToIdMap[name] = {scopeId, /*isStart=*/recordOp.getIsStart()}; + idToNameMap[scopeId] = name; + LDBG("Assigning new scope scopeId " << scopeId << " to op '" << recordOp << "'"); + opToIdMap[recordOp] = scopeId; + idToOpMap[scopeId] = recordOp; + scopeId++; + } else { + auto &[existingId, isStart] = nameToIdMap[name]; + if (isStart == recordOp.getIsStart()) { + // Error: duplicate start or end + mlir::emitError(recordOp.getLoc(), "The scope name '") + << name << "' has duplicate " + << (recordOp.getIsStart() ? "start" : "end") << " record"; + } else { + // Matching pair found + LDBG("Found matching pair for scope name '" << name << "' with scopeId " << existingId); + opToIdMap[recordOp] = existingId; + idToOpMap[existingId] = recordOp; + nameToIdMap.erase(name); + } + } + }); + + if (!nameToIdMap.empty()) { + for (auto &[name, idIsStartPair] : nameToIdMap) { + auto &[id, isStart] = idIsStartPair; + auto unclosedOp = idToOpMap.lookup(id); + mlir::emitError(unclosedOp.getLoc(), "The scope name '") + << name << "' is not properly closed (missing " + << (isStart ? "end" : "start") << " record)"; + } + } +} + void ScopeIdAllocation::reachability() { DenseMap inputBlockInfoMap; DenseMap outputBlockInfoMap; @@ -119,6 +168,7 @@ void ScopeIdAllocation::reachability() { virtualBlockList.emplace_back(block, Block::iterator()); }); + DenseSet exitVirtualBlocks; while (!virtualBlockList.empty()) { VirtualBlock &virtualBlock = virtualBlockList.front(); virtualBlockList.pop_front(); @@ -135,18 +185,17 @@ void ScopeIdAllocation::reachability() { break; } if (auto recordOp = dyn_cast(&op)) { - auto name = recordOp.getName(); - if (inputBlockInfo.contains(name)) { - if (!recordOp.getIsStart()) { - inputBlockInfo.erase(name); - } + auto scopeId = opToIdMap.lookup(recordOp); + if (recordOp.getIsStart()) { + inputBlockInfo.insert(scopeId); } else { - if (recordOp.getIsStart()) { - inputBlockInfo.insert(name); - } // else don't handle it right now as the scope might be monotonically closed later - } + inputBlockInfo.erase(scopeId); + } } } + if (successors.empty()) { + exitVirtualBlocks.insert(virtualBlock); + } // Get the reference because we want to update if it changed if (outputBlockInfoMap.count(virtualBlock) && inputBlockInfo == outputBlockInfoMap[virtualBlock]) { @@ -167,24 +216,21 @@ void ScopeIdAllocation::reachability() { for (auto iter : inputBlockInfoMap) { auto &virtualBlock = iter.first; auto inputBlockInfo = iter.second; - auto outputBlockInfo = outputBlockInfoMap[virtualBlock]; - DenseSet unclosedScopes; Block::iterator startIt = virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { if (auto recordOp = dyn_cast(&op)) { - auto name = recordOp.getName(); + auto scopeId = opToIdMap.lookup(recordOp); + auto name = idToNameMap.lookup(scopeId); if (recordOp.getIsStart()) { - if (inputBlockInfo.contains(name)) { + if (inputBlockInfo.contains(scopeId)) { mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is started without being closed"; } - inputBlockInfo.insert(name); - unclosedScopes.insert(name); + inputBlockInfo.insert(scopeId); } else { - if (inputBlockInfo.contains(name)) { - inputBlockInfo.erase(name); - unclosedScopes.erase(name); + if (inputBlockInfo.contains(scopeId)) { + inputBlockInfo.erase(scopeId); } else { mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is closed without being opened"; } @@ -194,50 +240,6 @@ void ScopeIdAllocation::reachability() { } } -void ScopeIdAllocation::liveness() { - // Stage 2: pair start/end records that refer to the same scope name and - // assign a numeric ID that downstream passes can reuse. - llvm::DenseMap> nameToIdMap; - llvm::DenseMap idToOpMap; - ScopeId scopeId = 0; - - funcOp->walk([&](RecordOp recordOp) { - auto name = recordOp.getName(); - LDBG("Processing RecordOp: " << recordOp); - if (!nameToIdMap.contains(name)) { - nameToIdMap[name] = {scopeId, /*isStart=*/recordOp.getIsStart()}; - idToNameMap[scopeId] = name; - LDBG("Assigning new scope scopeId " << scopeId << " to op '" << recordOp << "'"); - opToIdMap[recordOp] = scopeId; - idToOpMap[scopeId] = recordOp; - scopeId++; - } else { - auto &[existingId, isStart] = nameToIdMap[name]; - if (isStart == recordOp.getIsStart()) { - // Error: duplicate start or end - mlir::emitError(recordOp.getLoc(), "The scope name '") - << name << "' has duplicate " - << (recordOp.getIsStart() ? "start" : "end") << " record"; - } else { - // Matching pair found - LDBG("Found matching pair for scope name '" << name << "' with scopeId " << existingId); - opToIdMap[recordOp] = existingId; - idToOpMap[existingId] = recordOp; - nameToIdMap.erase(name); - } - } - }); - - if (!nameToIdMap.empty()) { - for (auto &[name, idIsStartPair] : nameToIdMap) { - auto &[id, isStart] = idIsStartPair; - auto unclosedOp = idToOpMap.lookup(id); - mlir::emitError(unclosedOp.getLoc(), "The scope name '") - << name << "' is not properly closed (missing " - << (isStart ? "end" : "start") << " record)"; - } - } -} void ScopeIdAllocation::dominance() { // Stage 3: determine parentage between scopes by checking dominance of start From 8450190b705bfbafba62734d2b76150b818c29b6 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Sat, 1 Nov 2025 20:54:53 -0400 Subject: [PATCH 30/33] Fix --- test/Proton/scope_id.mlir | 79 +++++----- test/lib/Proton/TestScopeIdAllocation.cpp | 6 +- .../lib/Analysis/ScopeIdAllocation.cpp | 139 +++++++++--------- .../proton/test/test_instrumentation.py | 11 +- 4 files changed, 113 insertions(+), 122 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 0e1afce6739c..ec2141338ef9 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -134,6 +134,44 @@ module { // ----- +module { + // expected-remark @below {{cf_if_only}} + tt.func @scf_cond(%cond: i1) { + scf.if %cond { + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} + proton.record start "if_only" + } + // expected-remark @below {{scope id = 0} + // expected-remark @below {{scope parent id = -1}} + proton.record end "if_only" + tt.return + } +} + +// ----- + +module { + tt.func @scf_loop_if(%cond: i1) { + %c0 = arith.constant 0 : index + scf.for %i = %c0 to %c0 step %c0 { + scf.if %cond { + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} + proton.record start "loop_if" + } + scf.if %cond { + // expected-remark @below {{scope id = 0}} + // expected-remark @below {{scope parent id = -1}} + proton.record end "loop_if" + } + } + tt.return + } +} + +// ----- + module { // expected-remark @below {{cf_single_branch}} tt.func @cf_single_branch(%cond: i1) { @@ -352,43 +390,4 @@ module { proton.record start "loop" cf.cond_br %cond, ^loop(%next : index), ^exit } -} - -// ----- - -module { - tt.func @cf_if_unclosed(%cond: i1) { - scf.if %cond { - // expected-error @below {{The scope name 'if_only' is not properly closed (missing end record)}} - proton.record start "if_only" - } - tt.return - } -} - -// ----- - -module { - tt.func @cf_duplicate_start() { - // expected-error @below {{The scope name 'dup_scope' is not properly closed (missing end record)}} - proton.record start "dup_scope" - // expected-error @below {{The scope name 'dup_scope' is started without being closed}} - // expected-error @below {{The scope name 'dup_scope' has duplicate start record}} - proton.record start "dup_scope" - tt.return - } -} - -// ----- - -module { - tt.func @cf_duplicate_end() { - // expected-error @below {{The scope name 'dup_scope' is closed without being opened}} - // expected-error @below {{The scope name 'dup_scope' is not properly closed (missing start record)}} - proton.record end "dup_scope" - // expected-error @below {{The scope name 'dup_scope' is closed without being opened}} - // expected-error @below {{The scope name 'dup_scope' has duplicate end record}} - proton.record end "dup_scope" - tt.return - } -} +} \ No newline at end of file diff --git a/test/lib/Proton/TestScopeIdAllocation.cpp b/test/lib/Proton/TestScopeIdAllocation.cpp index 8862d18f12c8..ef9ec57a4458 100644 --- a/test/lib/Proton/TestScopeIdAllocation.cpp +++ b/test/lib/Proton/TestScopeIdAllocation.cpp @@ -1,8 +1,6 @@ #include "mlir/Pass/Pass.h" #include "third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h" -#include - using namespace mlir; using namespace triton::proton; @@ -41,10 +39,10 @@ struct TestScopeIdAllocationPass funcOp.walk([&](RecordOp recordOp) { auto scopeId = moduleScopeIdAllocation.getOpScopeId(recordOp); mlir::emitRemark(recordOp.getLoc()) << "scope id = " << scopeId; - int64_t parentId = -1; + ScopeIdAllocation::ScopeId parentId = -1; if (auto parentIt = parentScopeIdMap.find(scopeId); parentIt != parentScopeIdMap.end()) - parentId = static_cast(parentIt->second); + parentId = parentIt->second; mlir::emitRemark(recordOp.getLoc()) << "scope parent id = " << parentId; }); diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index 738e124acf4a..a05c803c2bfa 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -50,63 +50,66 @@ struct BlockInfo { }; void ScopeIdAllocation::run() { - // Stage the analysis to match downstream consumers of scope metadata: + // We execute the following analysis stages in the order to verify if + // `proton.record` operations are well-formed and associate scope IDs for each + // pair of start/end records. // - // - liveness(): Pair start/end records and assign a shared numeric ID. - // There are multiple ways to pair start/end records: - // We choose a simple approach here: - // For each start record, we look for the nearest closing scope with the same name for pairing. + // 1. liveness() // - // Example - // MLIR: - // proton.record start @"foo" <- scopeId = 0 - // … - // proton.record end @"foo" <- scopeId = 0 - // … - // proton.record start @"foo" <- scopeId = 1 - // … - // proton.record end @"foo" <- scopeId = 1 + // Pair start/end records that share a name and assign a numeric + // identifier that later passes reuse. The current implementation pairs + // each start with the nearest matching end. // - // - reachability(): Track active scopes at CFG boundaries and flag malformed - // lifetimes. We optimistically collect all the "potentially" unclosed and - // closing scopes here, and validate them after the dataflow converges. + // proton.record start @"foo" // scopeId = 0 + // … + // proton.record end @"foo" // scopeId = 0 + // … + // proton.record start @"foo" // scopeId = 1 + // … + // proton.record end @"foo" // scopeId = 1 // - // Starting example MLIR: - // scf.if %cond { - // proton.record start @"foo" - // } - // Because `"foo"` never ends on the `then` branch, reachability() emits - // `The scope name 'foo' is not closed properly`. + // 2. reachability() // - // Valid example MLIR: - // scf.if %cond { - // proton.record start @"foo" - // } - // proton.record end @"foo" + // Track active scopes across CFG boundaries and surface + // malformed lifetimes once the dataflow converges. // - // We don't emit errors because we assume scf.if could be executed and it's - // up to the user to ensure proper semantics. + // scf.if %cond { + // proton.record start @"foo" + // } // - // - dominance(): - // (1) Check the dominance of start/end records to ensure well-formedness. - // Example MLIR: - // proton.record end @"foo" - // … - // proton.record start @"foo" + // Because `"foo"` never ends on the `then` branch, reachability() emits + // "The scope name 'foo' is not closed properly". // - // Because the end of `"foo"` dominates its start, dominance() emits an error. + // scf.if %cond { + // proton.record start @"foo" + // } + // proton.record end @"foo" // - // (2) Infer the parent/child hierarchy between scopes via - // dominance. Example MLIR: - // proton.record start @"outer" - // scf.if %cond { - // proton.record start @"inner" - // … - // proton.record end @"inner" - // } - // proton.record end @"outer" - // Because the start of `"outer"` dominates `"inner"`, dominance() records - // `(innerId -> outerId)` in `scopeParentIds`. + // No diagnostic is emitted: the pass assumes the branch may execute and + // leaves semantic responsibility to the caller. + // + // 3. dominance(): + // + // (a) Ensure that each start dominates its matching end. + // + // proton.record end @"foo" + // … + // proton.record start @"foo" + // + // Because the end dominates the start, dominance() reports an error. + // + // (b) Infer parent/child scope relationships using dominance facts. + // + // proton.record start @"outer" + // scf.if %cond { + // proton.record start @"inner" + // … + // proton.record end @"inner" + // } + // proton.record end @"outer" + // + // `"outer"` dominates `"inner"`, so dominance() records + // `(innerId -> outerId)` in `scopeParentIds`. liveness(); reachability(); dominance(); @@ -161,8 +164,7 @@ void ScopeIdAllocation::reachability() { std::deque virtualBlockList; funcOp->walk([&](Block *block) { - // Start the analysis from the entry blocks of any nested isolated from - // above regions. + // Seed the worklist with entry blocks of regions that are isolated-from-above. if (block->isEntryBlock() && !isa(block->getParentOp())) virtualBlockList.emplace_back(block, Block::iterator()); @@ -172,8 +174,7 @@ void ScopeIdAllocation::reachability() { while (!virtualBlockList.empty()) { VirtualBlock &virtualBlock = virtualBlockList.front(); virtualBlockList.pop_front(); - // Evaluate the transfer function for this block starting from the cached - // input state. + // Evaluate the transfer function for this block starting from the cached input state. auto inputBlockInfo = inputBlockInfoMap[virtualBlock]; SmallVector successors; Block::iterator startIt = @@ -190,29 +191,27 @@ void ScopeIdAllocation::reachability() { inputBlockInfo.insert(scopeId); } else { inputBlockInfo.erase(scopeId); - } + } } } if (successors.empty()) { exitVirtualBlocks.insert(virtualBlock); } - // Get the reference because we want to update if it changed + // Skip successor propagation if the output state is unchanged. if (outputBlockInfoMap.count(virtualBlock) && inputBlockInfo == outputBlockInfoMap[virtualBlock]) { - // If we have seen the block before and the inputBlockInfo is the same as - // the outputBlockInfo, we skip the successors continue; } - // Update the current block + // Update the current block. outputBlockInfoMap[virtualBlock].join(inputBlockInfo); - // Update the successors + // Propagate the new facts to successors. for (VirtualBlock &successor : successors) { inputBlockInfoMap[successor].join(outputBlockInfoMap[virtualBlock]); virtualBlockList.emplace_back(successor); } } - // Go through all blocks, validate reachability analysis results + // Validate the reachability analysis results for each block. for (auto iter : inputBlockInfoMap) { auto &virtualBlock = iter.first; auto inputBlockInfo = iter.second; @@ -242,8 +241,7 @@ void ScopeIdAllocation::reachability() { void ScopeIdAllocation::dominance() { - // Stage 3: determine parentage between scopes by checking dominance of start - // operations. + // Stage 3: derive scope parentage and verify dominance constraints. mlir::DominanceInfo domInfo(funcOp); llvm::DenseMap startRecordMap; llvm::DenseMap endRecordMap; @@ -272,17 +270,17 @@ void ScopeIdAllocation::dominance() { } auto sortedStartRecordOps = mlir::topologicalSort(startRecordOps); for (int i = 0; i < sortedStartRecordOps.size(); ++i) { - auto *op = sortedStartRecordOps[i]; - auto scopeId = opToIdMap.lookup(op); + auto *startOp = sortedStartRecordOps[i]; + auto scopeId = opToIdMap.lookup(startOp); auto endOp = endRecordMap.lookup(scopeId); for (int j = i - 1; j >= 0; --j) { auto *parentStartOp = sortedStartRecordOps[j]; auto parentScopeId = opToIdMap.lookup(parentStartOp); auto parentEndOp = endRecordMap.lookup(parentScopeId); - if (domInfo.dominates(parentStartOp, op) && + if (domInfo.dominates(parentStartOp, startOp) && domInfo.dominates(endOp, parentEndOp)) { auto parentId = opToIdMap.lookup(parentStartOp); - auto childId = opToIdMap.lookup(op); + auto childId = opToIdMap.lookup(startOp); scopeParentIds.push_back({childId, parentId}); break; } @@ -300,9 +298,8 @@ void ScopeIdAllocation::visitTerminator( } if (auto br = dyn_cast(op)) { - // The successors of an operation with regions can be queried via an - // interface. The operation branches to the entry blocks of its region - // successors. It can also branch to after itself. + // Query successors of an op-with-regions. The op can branch to region entry + // blocks or to the continuation after itself. SmallVector regions; br.getSuccessorRegions(RegionBranchPoint::parent(), regions); for (RegionSuccessor ®ion : regions) { @@ -320,8 +317,8 @@ void ScopeIdAllocation::visitTerminator( // reason. Check that the parent is actually a `RegionBranchOpInterface`. auto br = dyn_cast(op); if (br && isa(br->getParentOp())) { - // Check the successors of a region branch terminator. It can branch to - // another region of its parent operation or to after the parent op. + // Region branch terminators can jump to another region belonging to the + // parent operation or to the parent continuation. SmallVector operands(br->getNumOperands()); SmallVector regions; br.getSuccessorRegions(operands, regions); @@ -337,7 +334,7 @@ void ScopeIdAllocation::visitTerminator( return; } - // Otherwise, it could be a return op + // Otherwise, it could be a return-like op. if (op->hasTrait()) return; llvm_unreachable("Unknown terminator encountered in membar analysis"); diff --git a/third_party/proton/test/test_instrumentation.py b/third_party/proton/test/test_instrumentation.py index f00705748fd9..ee8e0839ee79 100644 --- a/third_party/proton/test/test_instrumentation.py +++ b/third_party/proton/test/test_instrumentation.py @@ -391,7 +391,6 @@ def matmul_kernel_tma(a_desc, b_desc, c_desc, # WARP_SPECIALIZE: tl.constexpr, # ): dtype = tl.float8e4nv if FP8_OUTPUT else tl.float16 - pl.enter_scope("kernel") pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) @@ -422,7 +421,6 @@ def matmul_kernel_tma(a_desc, b_desc, c_desc, # offs_cm = pid_m * BLOCK_SIZE_M offs_cn = pid_n * BLOCK_SIZE_N c_desc.store([offs_cm, offs_cn], c) - pl.exit_scope("kernel") def matmul_tma(a, b, warp_specialize: bool): # Check constraints. @@ -475,11 +473,10 @@ def grid(META): with open(temp_file, "rb") as f: data = json.load(f) - kernel_level = data[0]["children"][0]["children"][0] - assert kernel_level["children"][0]["frame"]["name"] == "loop" - assert kernel_level["children"][0]["metrics"]["cycles"] > 0 - assert kernel_level["frame"]["name"] == "kernel" - assert kernel_level["metrics"]["cycles"] > 0 + kernel = data[0]["children"][0] + assert kernel["children"][0]["frame"]["name"] == "loop" + assert kernel["children"][0]["metrics"]["cycles"] > 0 + assert kernel["frame"]["name"] == "matmul_kernel_tma" def test_timeline(tmp_path: pathlib.Path): From 52e9149cb3f182b517c9d7aaa5f0112e3d1e2875 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Sat, 1 Nov 2025 20:55:05 -0400 Subject: [PATCH 31/33] Fix --- test/Proton/scope_id.mlir | 2 +- test/lib/Proton/TestScopeIdAllocation.cpp | 3 +- .../include/Analysis/ScopeIdAllocation.h | 1 - .../lib/Analysis/ScopeIdAllocation.cpp | 47 ++++++++++--------- 4 files changed, 27 insertions(+), 26 deletions(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index ec2141338ef9..3ee96aa5ba44 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -390,4 +390,4 @@ module { proton.record start "loop" cf.cond_br %cond, ^loop(%next : index), ^exit } -} \ No newline at end of file +} diff --git a/test/lib/Proton/TestScopeIdAllocation.cpp b/test/lib/Proton/TestScopeIdAllocation.cpp index ef9ec57a4458..6a27b3f0af21 100644 --- a/test/lib/Proton/TestScopeIdAllocation.cpp +++ b/test/lib/Proton/TestScopeIdAllocation.cpp @@ -43,8 +43,7 @@ struct TestScopeIdAllocationPass if (auto parentIt = parentScopeIdMap.find(scopeId); parentIt != parentScopeIdMap.end()) parentId = parentIt->second; - mlir::emitRemark(recordOp.getLoc()) - << "scope parent id = " << parentId; + mlir::emitRemark(recordOp.getLoc()) << "scope parent id = " << parentId; }); }); } diff --git a/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h b/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h index 11bd9e57f9f7..ad06a3b3d6d5 100644 --- a/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h +++ b/third_party/proton/Dialect/include/Analysis/ScopeIdAllocation.h @@ -14,7 +14,6 @@ namespace mlir { namespace triton::proton { - class ScopeIdAllocation { public: using ScopeId = size_t; diff --git a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp index a05c803c2bfa..a39c2d17b914 100644 --- a/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp +++ b/third_party/proton/Dialect/lib/Analysis/ScopeIdAllocation.cpp @@ -1,5 +1,5 @@ -#include "mlir/Analysis/TopologicalSortUtils.h" #include "Analysis/ScopeIdAllocation.h" +#include "mlir/Analysis/TopologicalSortUtils.h" namespace mlir { namespace triton::proton { @@ -28,15 +28,11 @@ struct BlockInfo { return this->activeScopes.contains(scopeId); } - void erase(ScopeId scopeId) { - this->activeScopes.erase(scopeId); - } + void erase(ScopeId scopeId) { this->activeScopes.erase(scopeId); } - void insert(ScopeId scopeId) { - this->activeScopes.insert(scopeId); - } + void insert(ScopeId scopeId) { this->activeScopes.insert(scopeId); } - bool operator ==(const BlockInfo &other) const { + bool operator==(const BlockInfo &other) const { return this->activeScopes == other.activeScopes; } @@ -116,7 +112,8 @@ void ScopeIdAllocation::run() { } void ScopeIdAllocation::liveness() { - llvm::DenseMap> nameToIdMap; + llvm::DenseMap> + nameToIdMap; llvm::DenseMap idToOpMap; ScopeId scopeId = 0; @@ -126,7 +123,8 @@ void ScopeIdAllocation::liveness() { if (!nameToIdMap.contains(name)) { nameToIdMap[name] = {scopeId, /*isStart=*/recordOp.getIsStart()}; idToNameMap[scopeId] = name; - LDBG("Assigning new scope scopeId " << scopeId << " to op '" << recordOp << "'"); + LDBG("Assigning new scope scopeId " << scopeId << " to op '" << recordOp + << "'"); opToIdMap[recordOp] = scopeId; idToOpMap[scopeId] = recordOp; scopeId++; @@ -139,7 +137,8 @@ void ScopeIdAllocation::liveness() { << (recordOp.getIsStart() ? "start" : "end") << " record"; } else { // Matching pair found - LDBG("Found matching pair for scope name '" << name << "' with scopeId " << existingId); + LDBG("Found matching pair for scope name '" << name << "' with scopeId " + << existingId); opToIdMap[recordOp] = existingId; idToOpMap[existingId] = recordOp; nameToIdMap.erase(name); @@ -164,7 +163,8 @@ void ScopeIdAllocation::reachability() { std::deque virtualBlockList; funcOp->walk([&](Block *block) { - // Seed the worklist with entry blocks of regions that are isolated-from-above. + // Seed the worklist with entry blocks of regions that are + // isolated-from-above. if (block->isEntryBlock() && !isa(block->getParentOp())) virtualBlockList.emplace_back(block, Block::iterator()); @@ -174,11 +174,13 @@ void ScopeIdAllocation::reachability() { while (!virtualBlockList.empty()) { VirtualBlock &virtualBlock = virtualBlockList.front(); virtualBlockList.pop_front(); - // Evaluate the transfer function for this block starting from the cached input state. + // Evaluate the transfer function for this block starting from the cached + // input state. auto inputBlockInfo = inputBlockInfoMap[virtualBlock]; SmallVector successors; - Block::iterator startIt = - virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); + Block::iterator startIt = virtualBlock.second.isValid() + ? std::next(virtualBlock.second) + : virtualBlock.first->begin(); for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { if (op.hasTrait() || isa(op)) { @@ -215,8 +217,9 @@ void ScopeIdAllocation::reachability() { for (auto iter : inputBlockInfoMap) { auto &virtualBlock = iter.first; auto inputBlockInfo = iter.second; - Block::iterator startIt = - virtualBlock.second.isValid() ? std::next(virtualBlock.second) : virtualBlock.first->begin(); + Block::iterator startIt = virtualBlock.second.isValid() + ? std::next(virtualBlock.second) + : virtualBlock.first->begin(); for (Operation &op : llvm::make_range(startIt, virtualBlock.first->end())) { if (auto recordOp = dyn_cast(&op)) { auto scopeId = opToIdMap.lookup(recordOp); @@ -231,7 +234,8 @@ void ScopeIdAllocation::reachability() { if (inputBlockInfo.contains(scopeId)) { inputBlockInfo.erase(scopeId); } else { - mlir::emitError(recordOp.getLoc(), "The scope name '") << name << "' is closed without being opened"; + mlir::emitError(recordOp.getLoc(), "The scope name '") + << name << "' is closed without being opened"; } } } @@ -239,7 +243,6 @@ void ScopeIdAllocation::reachability() { } } - void ScopeIdAllocation::dominance() { // Stage 3: derive scope parentage and verify dominance constraints. mlir::DominanceInfo domInfo(funcOp); @@ -277,7 +280,7 @@ void ScopeIdAllocation::dominance() { auto *parentStartOp = sortedStartRecordOps[j]; auto parentScopeId = opToIdMap.lookup(parentStartOp); auto parentEndOp = endRecordMap.lookup(parentScopeId); - if (domInfo.dominates(parentStartOp, startOp) && + if (domInfo.dominates(parentStartOp, startOp) && domInfo.dominates(endOp, parentEndOp)) { auto parentId = opToIdMap.lookup(parentStartOp); auto childId = opToIdMap.lookup(startOp); @@ -288,8 +291,8 @@ void ScopeIdAllocation::dominance() { } } -void ScopeIdAllocation::visitTerminator( - Operation *op, SmallVector &successors) { +void ScopeIdAllocation::visitTerminator(Operation *op, + SmallVector &successors) { if (isa(op)) { // Collect the block successors of the branch. for (Block *successor : op->getSuccessors()) From 1aa1e53da4c8dac1ed91d7c093f32c6ec0db7d65 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Sat, 1 Nov 2025 21:10:27 -0400 Subject: [PATCH 32/33] Fix --- test/lib/Proton/TestScopeIdAllocation.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/lib/Proton/TestScopeIdAllocation.cpp b/test/lib/Proton/TestScopeIdAllocation.cpp index 6a27b3f0af21..888842a0f8a2 100644 --- a/test/lib/Proton/TestScopeIdAllocation.cpp +++ b/test/lib/Proton/TestScopeIdAllocation.cpp @@ -39,7 +39,7 @@ struct TestScopeIdAllocationPass funcOp.walk([&](RecordOp recordOp) { auto scopeId = moduleScopeIdAllocation.getOpScopeId(recordOp); mlir::emitRemark(recordOp.getLoc()) << "scope id = " << scopeId; - ScopeIdAllocation::ScopeId parentId = -1; + int64_t parentId = -1; if (auto parentIt = parentScopeIdMap.find(scopeId); parentIt != parentScopeIdMap.end()) parentId = parentIt->second; From 3c9ff174c82f333420ae62503480f55d205b6f33 Mon Sep 17 00:00:00 2001 From: Jokeren Date: Sat, 1 Nov 2025 21:16:56 -0400 Subject: [PATCH 33/33] Fix --- test/Proton/scope_id.mlir | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Proton/scope_id.mlir b/test/Proton/scope_id.mlir index 3ee96aa5ba44..2435e44f2c94 100644 --- a/test/Proton/scope_id.mlir +++ b/test/Proton/scope_id.mlir @@ -135,7 +135,7 @@ module { // ----- module { - // expected-remark @below {{cf_if_only}} + // expected-remark @below {{scf_cond}} tt.func @scf_cond(%cond: i1) { scf.if %cond { // expected-remark @below {{scope id = 0}}