Skip to content

[OpClassifier](fix) propagate per-result core_type for multi-result s… #1882

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
245516766 wants to merge 1 commit into
base: main-dev
Choose a base branch
from
+53 −48
Open

[OpClassifier](fix) propagate per-result core_type for multi-result s… #1882

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 53 additions & 48 deletions third_party/ascend/lib/DynamicCVPipeline/PlanComputeBlock/OpClassifier.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1313,12 +1313,8 @@ bool OpClassifierPass::handleYieldFromElseRegion(
}

// Get the core_type attribute from then region yield
auto thenCoreTypeAttr = thenYieldForElse->getAttr("ssbuffer.core_type");
if (!thenCoreTypeAttr) {
return false;
}

StringAttr thenCoreTypeStrAttr = dyn_cast<StringAttr>(thenCoreTypeAttr);
StringAttr thenCoreTypeStrAttr =
thenYieldForElse->getAttrOfType<StringAttr>(kCoreType);
if (!thenCoreTypeStrAttr) {
return false;
}
Expand Down Expand Up @@ -1411,11 +1407,45 @@ void OpClassifierPass::processYieldOperation(Operation *op,
// iter_arg that is passed through without modification), default to
// VECTOR since no compute op is associated with it.
if (Operation *def = operand.getDefiningOp()) {
OpCoreType ct = getCoreType(def);
if (ct == OP_CUBE_ONLY) {
coreTypes.push_back(OP_CUBE_ONLY);
// If def has multiple results, each result may have a different
// core_type.
// We need to find the specific result index and use the corresponding
// core_type from the stored multi-value attribute (e.g., "CUBE,VECTOR")
if (def->getNumResults() > 1) {
if (!isa<scf::SCFDialect>(def->getDialect())) {
LLVM_DEBUG(DBGS() << "[handleSCFYield] ERROR: def has multiple "
"results but is not scf dialect: "
<< *def << "\n");
return;
}
// Find which result index this operand corresponds to in def
unsigned resultIdx = 0;
for (unsigned idx = 0; idx < def->getNumResults(); ++idx) {
if (def->getResult(idx) == operand) {
resultIdx = idx;
break;
}
}

// Get core_type from def's attribute (stored as comma-separated
// multi-value)
if (auto ctStrAttr = def->getAttrOfType<StringAttr>(kCoreType)) {
std::string ctStr = ctStrAttr.getValue().str();
coreTypes.push_back(parseCoreTypeFromString(ctStr, resultIdx));
} else {
// Fallback: use getCoreType(def)
OpCoreType ct = getCoreType(def);
coreTypes.push_back(ct == OP_CUBE_ONLY ? OP_CUBE_ONLY
: OP_VECTOR_ONLY);
}
} else {
coreTypes.push_back(OP_VECTOR_ONLY);
// Non-scf operation: use getCoreType directly
OpCoreType ct = getCoreType(def);
if (ct == OP_CUBE_ONLY) {
coreTypes.push_back(OP_CUBE_ONLY);
} else {
coreTypes.push_back(OP_VECTOR_ONLY);
}
}
} else {
// BlockArgument (e.g., iter_arg passed through scf.for) -> VECTOR
Expand All @@ -1427,16 +1457,15 @@ void OpClassifierPass::processYieldOperation(Operation *op,
// Write ssbuffer.core_type attribute onto the yield op and its parent scf op
if (!coreTypes.empty()) {
std::string coreTypeStr = joinCoreTypes(coreTypes);
op->setAttr("ssbuffer.core_type",
StringAttr::get(op->getContext(), coreTypeStr));
op->setAttr(kCoreType, StringAttr::get(op->getContext(), coreTypeStr));

OpCoreType opCt = (coreTypeStr.find("CUBE") != std::string::npos)
? OP_CUBE_ONLY
: OP_VECTOR_ONLY;
opCoreTypes[op] = opCt;

if (parentOp && llvm::isa<scf::SCFDialect>(parentOp->getDialect())) {
parentOp->setAttr("ssbuffer.core_type",
parentOp->setAttr(kCoreType,
StringAttr::get(parentOp->getContext(), coreTypeStr));
opCoreTypes[parentOp] = opCt;
}
Expand All @@ -1457,44 +1486,20 @@ void OpClassifierPass::processYieldOperation(Operation *op,
// ssbuffer.core_type attribute reflecting the core_type of their result
// values (derived from the defining ops of each yield operand).
int OpClassifierPass::handleSCFYield() {
// Collect all scf.if operations
llvm::SmallVector<scf::IfOp> ifOps;
for (Operation *op : allOps) {
if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
ifOps.push_back(ifOp);
}
}

// Process scf.if operations: first then region, then else region
for (scf::IfOp ifOp : ifOps) {
Operation *thenYield = nullptr;
if (!ifOp.getThenRegion().empty()) {
thenYield = ifOp.getThenRegion().back().getTerminator();
}

if (isa<scf::YieldOp>(thenYield)) {
processYieldOperation(thenYield, nullptr);
}

if (!ifOp.getElseRegion().empty()) {
Operation *elseYield = ifOp.getElseRegion().back().getTerminator();
if (isa<scf::YieldOp>(elseYield)) {
processYieldOperation(elseYield, thenYield);
// Walk all scf.yield operations directly
getOperation().walk([&](scf::YieldOp yield) {

@olivervnc olivervnc Aug 31, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

walkWalkOrder::PostOrder

Operation *parentOp = yield->getParentOp();
Operation *thenYieldForElse = nullptr;

if (auto ifOp = dyn_cast<scf::IfOp>(parentOp)) {
if (yield->getParentRegion() != &ifOp.getThenRegion() &&
!ifOp.getThenRegion().empty()) {
thenYieldForElse = ifOp.getThenRegion().back().getTerminator();
}
}
}

// Process remaining scf.yield operations (not inside scf.if, e.g. scf.for)
for (Operation *op : allOps) {
if (!isa<scf::YieldOp>(op))
continue;

Operation *parentOp = op->getParentOp();
if (dyn_cast_or_null<scf::IfOp>(parentOp))
continue;

processYieldOperation(op, nullptr);
}
processYieldOperation(yield, thenYieldForElse);
});

return 0;
}
Expand Down
Loading