Skip to content

[mlir][Sol] Build precise per-type dispatch tables for internal fnptr calls #141

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
abinavpp wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][Sol] Build precise per-type dispatch tables for internal fnptr calls #141

Show file tree
Hide file tree
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
4 changes: 4 additions & 0 deletions mlir/include/mlir/Conversion/SolToYul/EVMUtil.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,10 @@ std::string clearStringTail();
/// the length word: `__sol.copy_string_data.<src>.<dst>`.
std::string copyStringData(mlir::sol::DataLocation src,
mlir::sol::DataLocation dst);
/// Dispatches an internal function pointer call over its candidate set:
/// `__sol.internal_dispatch.<rt|cr>.<id>...`.
std::string internalDispatch(bool runtime,
llvm::ArrayRef<int64_t> candidateIds);
} // namespace helpersym

/// IR Builder for EVM specific lowering.
Expand Down
3 changes: 3 additions & 0 deletions mlir/include/mlir/Dialect/Sol/SolBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,9 @@ def Sol_FuncRefType : Sol_Type<"FuncRef", "func_ref"> {
let assemblyFormat = [{ `<` $funcTy `>` }];
}

// A func ref, or the i256 function id it legalizes to.
def Sol_FuncRefLike : AnyTypeOf<[Sol_FuncRefType, I256]>;

// This represents an external function reference (address + selector).
def Sol_ExtFuncRefType : Sol_Type<"ExtFuncRef", "ext_func_ref"> {
let parameters = (ins "FunctionType":$funcTy);
Expand Down
18 changes: 10 additions & 8 deletions mlir/include/mlir/Dialect/Sol/SolOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -788,7 +788,7 @@ def Sol_DefaultStorageOp : Sol_Op<"default_storage", [Pure]> {
}

def Sol_DefaultFuncConstantOp : Sol_Op<"default_func_constant", [Pure]> {
let results = (outs AnyType:$addr);
let results = (outs Sol_FuncRefType:$addr);

let builders = [
OpBuilder<(ins), [{
Expand All @@ -804,7 +804,7 @@ def Sol_FuncConstantOp : Sol_Op<"func_constant", [Pure]> {
// TODO: SymbolUserOpInterface verifier: $sym must resolve to a sol.func with
// an assigned id.
let arguments = (ins FlatSymbolRefAttr:$sym);
let results = (outs AnyType:$addr);
let results = (outs Sol_FuncRefLike:$addr);

let assemblyFormat = "$sym attr-dict `:` type($addr)";
}
Expand All @@ -831,10 +831,10 @@ def Sol_ExtFuncSelectorOp : Sol_Op<"ext_func_selector", [Pure]> {
}

// Base class for indirect call ops (internal and external function pointers).
class Sol_ICallOpBase<string mnemonic, dag extraArgs = (ins),
string extraAsmFmt = "">
class Sol_ICallOpBase<string mnemonic, TypeConstraint calleeTy,
dag extraArgs = (ins), string extraAsmFmt = "">
: Sol_Op<mnemonic, [CallOpInterface]> {
let arguments = !con((ins AnyType:$callee, Variadic<AnyType>:$callee_operands),
let arguments = !con((ins calleeTy:$callee, Variadic<AnyType>:$callee_operands),
extraArgs,
(ins OptionalAttr<DictArrayAttr>:$arg_attrs,
OptionalAttr<DictArrayAttr>:$res_attrs));
Expand Down Expand Up @@ -865,11 +865,13 @@ class Sol_ICallOpBase<string mnemonic, dag extraArgs = (ins),
"attr-dict `:` type($callee) `,` functional-type($callee_operands, results)");
}

// Indirect call for internal function pointers.
def ICallOp : Sol_ICallOpBase<"icall">;
// Indirect call for internal function pointers. $candidates holds the ids of
// the functions referenced by a sol.func_constant of the callee type
def ICallOp : Sol_ICallOpBase<"icall", Sol_FuncRefLike,
(ins OptionalAttr<DenseI64ArrayAttr>:$candidates)>;

// Indirect call for external function pointers.
def ExtICallOp : Sol_ICallOpBase<"ext_icall",
def ExtICallOp : Sol_ICallOpBase<"ext_icall", Sol_ExtFuncRefType,
(ins UI256:$gas, UI256:$value, UnitAttr:$static_call, UnitAttr:$try_call),
"`gas` $gas `value` $value ">;

Expand Down
11 changes: 11 additions & 0 deletions mlir/lib/Conversion/SolToYul/EVMUtil.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2171,6 +2171,17 @@ std::string evm::helpersym::copyStringData(sol::DataLocation src,
sol::stringifyDataLocation(dst).lower()});
}

std::string evm::helpersym::internalDispatch(bool runtime,
ArrayRef<int64_t> candidateIds) {
SmallVector<std::string> idStrs;
idStrs.reserve(candidateIds.size());
for (int64_t id : candidateIds)
idStrs.push_back(std::to_string(id));
SmallVector<StringRef> fields{runtime ? "rt" : "cr"};
llvm::append_range(fields, idStrs);
return makeHelperSymbol("internal_dispatch", fields);
}

sol::FuncOp evm::Builder::getOrCreateHelperFn(
StringRef symbol, TypeRange argTys, TypeRange resTys,
llvm::function_ref<void(ValueRange)> genBody, Location loc) {
Expand Down
109 changes: 85 additions & 24 deletions mlir/lib/Conversion/SolToYul/SolToYul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3952,39 +3952,64 @@ struct ICallOpLowering : public OpConversionPattern<sol::ICallOp> {
// is lowered, so the dispatch table can still inspect the Sol symbol table.
using OpConversionPattern<sol::ICallOp>::OpConversionPattern;

LogicalResult matchAndRewrite(sol::ICallOp op, OpAdaptor adaptor,
ConversionPatternRewriter &r) const override {
/// Generates the dispatch function that switches over the candidate
/// functions of the caller's phase. Outlining this reduces the caller's
/// stack pressure.
yul::FuncOp genDispatchFn(StringRef name, ArrayRef<int64_t> candidateIds,
FunctionType calleeTy, Type calleeIdTy,
bool callerRuntime, sol::ICallOp op,
ConversionPatternRewriter &r) const {
Location loc = op.getLoc();
Operation *symTab = SymbolTable::getNearestSymbolTable(op);
evm::Builder evmB(getModule(op), r, loc);

auto calleeArgs = adaptor.getOperands().drop_front();
SmallVector<Type> convertedResTys;
if (failed(getTypeConverter()->convertTypes(op.getResultTypes(),
convertedResTys)))
return failure();

auto calleeTy = mlir::FunctionType::get(
r.getContext(), calleeArgs.getTypes(), convertedResTys);

// Collect functions with matching signature.
Operation *symTab = SymbolTable::getNearestSymbolTable(op);
bool callerRuntime = op->getParentOfType<sol::FuncOp>().getRuntime();
// Resolve the candidate ids in the caller's phase.
DenseSet<int64_t> candidateIdSet(candidateIds.begin(), candidateIds.end());
SmallVector<int64_t> caseIds;
SmallVector<sol::FuncOp> caseFns;
symTab->walk([&](sol::FuncOp fn) {
// At this stage the nearest Sol symbol table can still contain both
// creation and runtime functions. Internal function pointers are only
// valid within the caller's phase, so do not dispatch across that
// boundary just because the id and signature match.
// boundary just because the id matches.
if (fn.getRuntime() != callerRuntime)
return;
if (fn.getId() && fn.getFunctionType() == calleeTy) {
if (fn.getId() && candidateIdSet.contains(*fn.getId())) {
assert(fn.getFunctionType() == calleeTy &&
"candidate signature mismatch");
caseFns.push_back(fn);
caseIds.push_back(*fn.getId());
}
});

auto calleeIntTy = cast<IntegerType>(adaptor.getCallee().getType());
OpBuilder::InsertionGuard insertGuard(r);
r.setInsertionPointAfter(op->getParentOfType<sol::FuncOp>());

SmallVector<Type> inputTys{calleeIdTy};
llvm::append_range(inputTys, calleeTy.getInputs());
auto fnTy = mlir::FunctionType::get(r.getContext(), inputTys,
calleeTy.getResults());
auto fn = r.create<yul::FuncOp>(loc, name, fnTy);

SmallVector<Location> argLocs(inputTys.size(), loc);
Block *entry =
r.createBlock(&fn.getBody(), fn.getBody().end(), inputTys, argLocs);
r.setInsertionPointToStart(entry);
ValueRange calleeArgs = entry->getArguments().drop_front();

// No function of this signature is ever taken as a value, so no valid id
// can reach this call: unconditionally panic.
if (caseFns.empty()) {
evmB.genPanic(mlir::evm::PanicCode::InvalidInternalFunction);
SmallVector<Value> undefs;
undefs.reserve(calleeTy.getNumResults());
for (Type ty : calleeTy.getResults())
undefs.push_back(r.create<LLVM::UndefOp>(loc, ty));
r.create<yul::FuncReturnOp>(loc, undefs);
return fn;
}

auto calleeIntTy = cast<IntegerType>(calleeIdTy);
SmallVector<APInt> caseVals;
caseVals.reserve(caseIds.size());
for (int64_t caseId : caseIds)
Expand All @@ -3993,24 +4018,60 @@ struct ICallOpLowering : public OpConversionPattern<sol::ICallOp> {
RankedTensorType::get(static_cast<int64_t>(caseVals.size()),
calleeIntTy),
caseVals);
auto switchOp = r.create<yul::SwitchOp>(
loc, convertedResTys, adaptor.getCallee(), caseIdsAttr, caseIds.size());
auto switchOp = r.create<yul::SwitchOp>(loc, calleeTy.getResults(),
entry->getArgument(0), caseIdsAttr,
caseIds.size());
for (size_t i = 0; i < caseFns.size(); ++i) {
r.setInsertionPointToStart(&switchOp.getCaseRegions()[i].emplaceBlock());
auto call = r.create<yul::FuncCallOp>(
loc, convertedResTys, FlatSymbolRefAttr::get(caseFns[i]), calleeArgs);
auto call = r.create<yul::FuncCallOp>(loc, calleeTy.getResults(),
FlatSymbolRefAttr::get(caseFns[i]),
calleeArgs);
r.create<yul::YieldOp>(loc, call.getResults());
}

r.setInsertionPointToStart(&switchOp.getDefaultRegion().emplaceBlock());
evmB.genPanic(mlir::evm::PanicCode::InvalidInternalFunction);
SmallVector<Value> undefs;
undefs.reserve(op.getNumResults());
for (Type ty : convertedResTys)
undefs.reserve(calleeTy.getNumResults());
for (Type ty : calleeTy.getResults())
undefs.push_back(r.create<LLVM::UndefOp>(loc, ty));
r.create<yul::YieldOp>(loc, undefs);

r.replaceOp(op, switchOp.getResults());
r.setInsertionPointToEnd(entry);
r.create<yul::FuncReturnOp>(loc, switchOp.getResults());
return fn;
}

LogicalResult matchAndRewrite(sol::ICallOp op, OpAdaptor adaptor,
ConversionPatternRewriter &r) const override {
auto calleeArgs = adaptor.getOperands().drop_front();

@PavelKopyl PavelKopyl Aug 14, 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.

I suggest to add an assert that the enclosing function's ancestor is already a yul.object.

SmallVector<Type> convertedResTys;
if (failed(getTypeConverter()->convertTypes(op.getResultTypes(),
convertedResTys)))
return failure();

auto calleeTy = mlir::FunctionType::get(
r.getContext(), calleeArgs.getTypes(), convertedResTys);

// Set at pass entry, before the sol signatures are legalized.
std::optional<ArrayRef<int64_t>> candidateIds = op.getCandidates();
assert(candidateIds && "sol.icall without candidate annotation");
Comment on lines +4057 to +4058

@PavelKopyl PavelKopyl Aug 14, 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.

It looks like we have here an API mismatch: in TableGen candidates are declared as optional, OptionalAttr<DenseI64ArrayAttr>:$candidates meaning a sol.icall without this attribute is a fully valid op, but at this stage (to be more specific, after annotateICallCandidates) we require it's presence. We could introduce another opcode, ResolvedICallOp where this attribute is not optional. annotateICallCandidates would convert CallOp to the new operation and declaring CallOp invalid.


Operation *symTab = SymbolTable::getNearestSymbolTable(op);
bool callerRuntime = op->getParentOfType<sol::FuncOp>().getRuntime();
std::string dispatchFnName =
evm::helpersym::internalDispatch(callerRuntime, *candidateIds);

@PavelKopyl PavelKopyl Aug 14, 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.

It looks like we may get a symbol collision for empty candidateIds. For example:

Contract Test {
function (uint256) internal returns (uint256) f1;
function () internal f2;
function a(uint256 x) public returns (uint256) { return f1(x); }
function b() public { f2(); }
}

causes a crash, as it produces one func.func @__sol.internal_dispatch.rt : (i256, i256) -> i256 and two incompatible call sites. I'd suggest to inline an empty dispatch function (that contains only Panic(0x51) call).

auto dispatchFn = dyn_cast_or_null<yul::FuncOp>(
SymbolTable::lookupSymbolIn(symTab, dispatchFnName));
Comment on lines +4064 to +4065

@PavelKopyl PavelKopyl Aug 14, 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.

I'd suggest to add here more checks that we do not cause symbol clashing with:

  • non functions
  • functions, but with different signatures
Operation *found = SymbolTable::lookupSymbolIn(symTab, dispatchFnName);
auto dispatchFn = dyn_cast_or_null<yul::FuncOp>(found);
assert((!found || dispatchFn) && "dispatch symbol clashes with non-function");
assert(!dispatchFn || dispatchFn.getFunctionType() == expectedDispatchTy);

if (!dispatchFn)
dispatchFn =
genDispatchFn(dispatchFnName, *candidateIds, calleeTy,
adaptor.getCallee().getType(), callerRuntime, op, r);

SmallVector<Value> callArgs{adaptor.getCallee()};
llvm::append_range(callArgs, calleeArgs);
r.replaceOpWithNewOp<yul::FuncCallOp>(
op, convertedResTys, FlatSymbolRefAttr::get(dispatchFn), callArgs);
return success();
}
};
Expand Down
37 changes: 37 additions & 0 deletions mlir/lib/Conversion/SolToYul/SolToYulPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "mlir/Dialect/Sol/Sol.h"
#include "mlir/Dialect/Yul/Yul.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/SetVector.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTSOLTOYULPASS
Expand Down Expand Up @@ -295,9 +296,45 @@ struct ConvertSolToYulPass
return success();
}

/// Sets the candidates attribute of every sol.icall.
void annotateICallCandidates(ModuleOp mod) {
Builder b(mod.getContext());
DenseMap<std::pair<Operation *, sol::FuncRefType>, SetVector<int64_t>>
buckets;
SmallVector<sol::ICallOp> icalls;
mod.walk([&](Operation *op) {
if (auto icall = dyn_cast<sol::ICallOp>(op)) {
icalls.push_back(icall);
return;
}
auto fnConst = dyn_cast<sol::FuncConstantOp>(op);
if (!fnConst)
return;
auto fn = SymbolTable::lookupNearestSymbolFrom<sol::FuncOp>(
fnConst, fnConst.getSymAttr());
assert(fn && fn.getId() && "unresolvable sol.func_constant");
if (!fn || !fn.getId())
return;
Comment on lines +315 to +317

@PavelKopyl PavelKopyl Aug 14, 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.

The assert and the consequent if statement condition contradict to each other. I think we could use an llvm_unreachable here.

buckets[{fnConst->getParentOfType<ModuleOp>().getOperation(),
cast<sol::FuncRefType>(fnConst.getType())}]
.insert(*fn.getId());
});
for (sol::ICallOp icall : icalls) {
SmallVector<int64_t> ids;
auto it =
buckets.find({icall->getParentOfType<ModuleOp>().getOperation(),
cast<sol::FuncRefType>(icall.getCallee().getType())});
if (it != buckets.end())
ids.assign(it->second.begin(), it->second.end());
llvm::sort(ids);

@PavelKopyl PavelKopyl Aug 14, 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.

I suggest to move sorting from this loop to do sort only once, finalize sorted attribute creation and then use them to annotate icalls. We could do something like:


  DenseMap<std::pair<Operation *, sol::FuncRefType>, DenseI64ArrayAttr>
      bucketAttrs;
  for (auto &[key, idSet] : buckets) {
    SmallVector<int64_t> ids = idSet.takeVector();
    llvm::sort(ids);
    bucketAttrs[key] = b.getDenseI64ArrayAttr(ids);
  }

  DenseI64ArrayAttr emptyAttr = b.getDenseI64ArrayAttr({});
  for (sol::ICallOp icall : icalls) {
    auto it =
        bucketAttrs.find({icall->getParentOfType<ModuleOp>().getOperation(),
                          cast<sol::FuncRefType>(icall.getCallee().getType())});
    icall.setCandidatesAttr(it != bucketAttrs.end() ? it->second : emptyAttr);
  }

icall.setCandidatesAttr(b.getDenseI64ArrayAttr(ids));
}
}

void runOnOperation() override {
ModuleOp mod = getOperation();
evm::SolTypeConverter tyConv;
annotateICallCandidates(mod);
if (failed(runStage1Conversion(mod, tyConv))) {
signalPassFailure();
return;
Expand Down