Skip to content

reactantXLAExec: address an offset argument view as a slice of its base buffer #3286

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

Merged
wsmoses merged 1 commit into from
Sep 15, 2026
Merged

reactantXLAExec: address an offset argument view as a slice of its base buffer #3286

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
58 changes: 52 additions & 6 deletions deps/ReactantExtra/API.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3645,19 +3645,20 @@ REACTANT_ABI void reactantXLAExec(LinkableRuntime **__restrict__ lrtP,

std::vector<std::vector<int64_t>> sizeKey;
sizeKey.reserve(argcnt + (constcnt ? 1 : 0));
// An argument may point into the middle of an allocation (a block
// operator's sub-vector view). The executable then takes the whole base
// buffer and addresses the view as a slice of it; the offsets select the
// executable, so they are part of the cache key.
std::vector<int64_t> viewOffset(argcnt, 0);
for (int64_t i = 0; i < argcnt; i++) {
auto &&[argB, argO, argP] = bufferAndOffset(lrt, args[i]);
if (argO != 0) {
llvm::errs() << "only zero-offset execution supported, argument " << i
<< " had byte offset of " << argO << " (ptr=" << args[i]
<< ", resolved base=" << (void *)argP << ")\n";
exit(1);
}
viewOffset[i] = argO;
baseArrays[i] = argB;
basePtrs[i] = argP;
auto dims = argB->on_device_shape().dimensions();
sizeKey.emplace_back(dims.begin(), dims.end());
}
sizeKey.push_back(viewOffset);

// The specialized scalars are part of what the executable was compiled
// for, so they are part of what it is cached under.
Expand Down Expand Up @@ -3730,6 +3731,51 @@ REACTANT_ABI void reactantXLAExec(LinkableRuntime **__restrict__ lrtP,
baseArrays[i]->on_device_shape(), builder));
types.push_back(RTT);
}
// A view argument becomes its base buffer: every use of it reads the
// static slice at the view's offset, and its result is written back into
// the base at that offset. The optimizer folds the update of an unchanged
// view back to the base argument, so a read-only view is recognized by
// identity like any other argument.
{
auto ret = cast<func::ReturnOp>(funcOp.getBody().front().back());
for (int64_t i = 0; i < argcnt; i++) {
if (!viewOffset[i])
continue;
auto baseTy = cast<mlir::RankedTensorType>(types[i]);
int64_t elemBytes = baseTy.getElementTypeBitWidth() / 8;
if (baseTy.getRank() != 1 || elemBytes == 0 ||
viewOffset[i] % elemBytes) {
llvm::errs() << " offset view of argument " << i << " (byte offset "
<< viewOffset[i]
<< ") does not slice its buffer of type " << baseTy
<< "\n";
exit(1);
}
int64_t start = viewOffset[i] / elemBytes;
auto arg = funcOp.getArgument(i);
arg.setType(baseTy);
mlir::OpBuilder b(&funcOp.getBody().front(),
funcOp.getBody().front().begin());
auto slice = b.create<mlir::stablehlo::SliceOp>(
funcOp.getLoc(), arg, b.getDenseI64ArrayAttr({start}),
b.getDenseI64ArrayAttr({baseTy.getDimSize(0)}),
b.getDenseI64ArrayAttr({1}));
arg.replaceAllUsesExcept(slice, slice);
mlir::OpBuilder rb(ret);
auto idxTy = mlir::RankedTensorType::get({}, rb.getI64Type());
auto startCst = rb.create<mlir::stablehlo::ConstantOp>(
funcOp.getLoc(), mlir::DenseElementsAttr::get(idxTy, start));
auto dus = rb.create<mlir::stablehlo::DynamicUpdateSliceOp>(
funcOp.getLoc(), arg, ret.getOperand(i),
mlir::ValueRange{startCst});
ret->setOperand(i, dus);
}
// Only the argument types change here; the result types are left to
// the refinement passes, which pin them from the arguments.
funcOp.setType(mlir::FunctionType::get(
module->getContext(), funcOp.getBody().front().getArgumentTypes(),
funcOp.getFunctionType().getResults()));
}
if (constcnt) {
// The injected constants make the dynamic shape operands static;
// statify the module before argument refinement so the result types
Expand Down
Loading