当前位置：首页 > news >正文

mlir operand

news 2025/8/15 10:44:52

std::pair<unsigned, unsigned> getODSOperandIndexAndLength(unsigned idx) 函数

它是 MLIR 里 ODS（Operation Definition Specification）自动生成。
idx 是算子的第idx个operand在 ODS operand td文件中的声明里的顺序编号，从 0 开始
比如你这个 Sch_KernelOp 的td文件中 operand 是这样排的：

Variadic<AnyTypeOf<...>>:$watchOperands,    // 段 0
Variadic<Sch_SyncCounterType>:$advanceOperands, // 段 1
Variadic<AnyRankedTensor>:$inits           // 段 2

因为operand的类型可以是Variadic，所以这个函数需要返回起始位置和长度
getODSOperandIndexAndLength(1) → 返回 advanceOperands 在总operands里的起始位置, 长度

// 假设 kernelOp 是 sch::KgraphKernelOp
auto idxAndLen = kernelOp.getODSOperandIndexAndLength(1); // 段 0
unsigned startIdx = idxAndLen.first;   // operand(1)在总Operands里的起点
unsigned length   = idxAndLen.second;  // operand(1)在总Operands里的长度llvm::outs() << "watchOperands start at " << startIdx << " length = " << length << "\n";// 获取这一段的 operands
auto operands = kernelOp->getOperands().slice(startIdx, length);

getOperands

获取所有的operands

setOperands

set一部分operands

SmallVector<Value> newWatchauto idxAndLen = kernelOp.getODSOperandIndexAndLength(1); // 段 1
unsigned startIdx = idxAndLen.first;   // 在 getOperands() 里的起点
unsigned length   = idxAndLen.second;  // 段的长度// 替换operand(1)这一段 operands
kernelOp->setOperands(startIdx, length, newWatch);

set所有的operands

// 构造新的 watchOperands：加上新 kernel 的 token
SmallVector<Value> fullOperands;
fullOperands.push_back(newKernel.getResult(0));                      // 新 watch token
fullOperands.append(kernelOp.getWatchOperands().begin(),kernelOp.getWatchOperands().end());              // 原 watch
fullOperands.append(kernelOp.getAdvanceOperands().begin(),kernelOp.getAdvanceOperands().end());            // advance
fullOperands.append(kernelOp.getInits().begin(),kernelOp.getInits().end());                      // inits// 一次性替换所有操作数
kernelOp->setOperands(fullOperands);

注：修改了operand之后，如果涉及某个variadic operand里的个数变量，则需要修改"operandSegmentSizes"属性里对应operand的值

if (transposeParentKernelOp->hasAttr("operandSegmentSizes")) {auto segVals = cast<DenseI32ArrayAttr>(transposeParentKernelOp->getAttr("operandSegmentSizes")).asArrayRef().vec();auto segSize = segVals.size();if (segSize > 0) {segVals[0] = 1;}auto segAttr = DenseI32ArrayAttr::get(transposeParentKernelOp.getContext(), segVals);transposeParentKernelOp->setAttr("operandSegmentSizes", segAttr);
}

setOperand(index)

修改op的某个operand

注：如果修改了operand的类型，那么block argument的类型也要修改

// Update region block argument types
Region &bodyRegion = linalgOp->getRegion(0);
if (!bodyRegion.empty()) {Block &bodyBlock = bodyRegion.front();if (bodyBlock.getNumArguments() >= 2) {// Update argument types to match swapped operandsbodyBlock.getArgument(0).setType(input1Type.getElementType()); // 新操作数0类型bodyBlock.getArgument(1).setType(input0Type.getElementType()); // 新操作数1类型}
}