Skip to content

Commit 25d01e0

Browse files
Store function signatures instead of FuncOp handles in functionMap (review doc 2)
NamespaceInfo::functionMap cached raw mlir_ts::FuncOp handles; since the discovery pass emits into a throwaway module (#213), any handle registered during discovery is guaranteed to dangle once the module is erased - consumers survived only because they read the (context-owned) type early. The map now stores FunctionEntry { symbol name, function type }, which is all any consumer used; the one site that needs a live op (generic-instantiation short-circuit) resolves it through theModule.lookupSymbol with a short/full-name fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent f0bf983 commit 25d01e0

2 files changed

Lines changed: 32 additions & 12 deletions

File tree

tslang/include/TypeScript/MLIRLogic/MLIRGenStore.h

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -985,6 +985,20 @@ struct GenericClassInfo
985985
}
986986
};
987987

988+
// what we know about a registered function; deliberately not the FuncOp itself -
989+
// discovery-pass ops are erased with the discovery module, so cached op handles dangle.
990+
// Resolve a live op through theModule.lookupSymbol when one is actually needed.
991+
struct FunctionEntry
992+
{
993+
std::string name;
994+
mlir_ts::FunctionType funcType;
995+
996+
explicit operator bool() const
997+
{
998+
return static_cast<bool>(funcType);
999+
}
1000+
};
1001+
9881002
struct NamespaceInfo
9891003
{
9901004
public:
@@ -998,7 +1012,7 @@ struct NamespaceInfo
9981012

9991013
llvm::StringMap<mlir_ts::FunctionType> functionTypeMap;
10001014

1001-
llvm::StringMap<mlir_ts::FuncOp> functionMap;
1015+
llvm::StringMap<FunctionEntry> functionMap;
10021016

10031017
llvm::StringMap<GenericFunctionInfo::TypePtr> genericFunctionMap;
10041018

tslang/lib/TypeScript/MLIRGen.cpp

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5312,7 +5312,7 @@ class MLIRGenImpl
53125312
auto funcIt = getFunctionMap().find(name);
53135313
if (funcIt != getFunctionMap().end())
53145314
{
5315-
auto cachedFuncType = funcIt->second.getFunctionType();
5315+
auto cachedFuncType = funcIt->second.funcType;
53165316
if (cachedFuncType.getNumResults() > 0)
53175317
{
53185318
auto returnType = cachedFuncType.getResult(0);
@@ -6032,7 +6032,8 @@ class MLIRGenImpl
60326032
auto name = funcProto->getNameWithoutNamespace();
60336033
if (!getFunctionMap().count(name))
60346034
{
6035-
getFunctionMap().insert({name, funcOp});
6035+
getFunctionMap().insert(
6036+
{name, FunctionEntry{funcOp.getName().str(), mlir::cast<mlir_ts::FunctionType>(funcOp.getFunctionType())}});
60366037

60376038
LLVM_DEBUG(llvm::dbgs() << "\n!! reg. func: " << name << " type:" << funcOp.getFunctionType() << " function name: " << funcProto->getName()
60386039
<< " num inputs:" << mlir::cast<mlir_ts::FunctionType>(funcOp.getFunctionType()).getNumInputs()
@@ -6081,10 +6082,18 @@ class MLIRGenImpl
60816082
{
60826083
auto [fullFunctionName, functionName] = getNameOfFunction(functionLikeDeclarationBaseAST, funcDeclGenContext);
60836084

6084-
auto funcOp = lookupFunctionMap(functionName);
6085-
if (funcOp && theModule.lookupSymbol(functionName)
6085+
auto funcEntry = lookupFunctionMap(functionName);
6086+
if (funcEntry && theModule.lookupSymbol(functionName)
60866087
|| theModule.lookupSymbol(fullFunctionName))
60876088
{
6089+
// resolve a live op from the module instead of returning a cached handle;
6090+
// the registered symbol is usually the full name
6091+
auto funcOp = theModule.lookupSymbol<mlir_ts::FuncOp>(functionName);
6092+
if (!funcOp)
6093+
{
6094+
funcOp = theModule.lookupSymbol<mlir_ts::FuncOp>(fullFunctionName);
6095+
}
6096+
60886097
return {mlir::success(), funcOp, functionName, false};
60896098
}
60906099
}
@@ -15966,11 +15975,8 @@ class MLIRGenImpl
1596615975
auto fn = getFunctionMap().find(name);
1596715976
if (fn != getFunctionMap().end())
1596815977
{
15969-
auto funcOp = fn->getValue();
15970-
auto funcType = funcOp.getFunctionType();
15971-
auto funcName = funcOp.getName();
15972-
15973-
return resolveFunctionWithCapture(location, funcName, funcType, mlir::Value(), false, genContext);
15978+
auto &funcEntry = fn->getValue();
15979+
return resolveFunctionWithCapture(location, funcEntry.name, funcEntry.funcType, mlir::Value(), false, genContext);
1597415980
}
1597515981

1597615982
return mlir::Value();
@@ -26273,12 +26279,12 @@ genContext);
2627326279
lookupLogic(functionTypeMap);
2627426280
}
2627526281

26276-
auto getFunctionMap() -> llvm::StringMap<mlir_ts::FuncOp> &
26282+
auto getFunctionMap() -> llvm::StringMap<FunctionEntry> &
2627726283
{
2627826284
return currentNamespace->functionMap;
2627926285
}
2628026286

26281-
auto lookupFunctionMap(StringRef name) -> mlir_ts::FuncOp
26287+
auto lookupFunctionMap(StringRef name) -> FunctionEntry
2628226288
{
2628326289
lookupLogic(functionMap);
2628426290
}

0 commit comments

Comments
 (0)