Files
scriptc/native/llvm-codegen/src/partition.cpp
T
Chris Tate 0e67cd1e3b Place functions too large to inline by their own name
The helper keeps an internal function that exactly one other function
references in its referrer's partition, so the partition can inline it.
A program's module initializers are each called only from the entry
point, so all of them, with every single-use helper they call, shared
the entry point's partition. In the three.js nodes fixture that
partition held 1.1 million lines of LLVM IR, almost five times the
average, and set the native build's longest compile.

A function whose body is larger than 256 KiB of LLVM text is past
LLVM's inline cost limit even for the only call to a local function, so
it is now placed by the hash of its own name. Its own single-use
callees still stay with it.
2026-10-18 08:02:34 -05:00

590 lines
19 KiB
C++

#include "partition.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/ThreadPool.h"
#include "llvm/Support/Threading.h"
#include <algorithm>
#include <cstdint>
using namespace llvm;
namespace scriptc {
namespace {
constexpr unsigned NoPartition = ~0u;
enum class SymbolKind : uint8_t { Function, Variable, Declaration };
struct Symbol {
SymbolKind Kind = SymbolKind::Declaration;
bool Local = false;
StringRef Name;
// Definition header line, or the complete declaration line.
StringRef Header;
// Function lines after the header, through the closing brace.
StringRef Body;
// Declaration other partitions use for this definition.
std::string Declaration;
// Local definitions are rewritten with hidden visibility.
std::string PromotedHeader;
std::vector<unsigned> Refs;
std::vector<unsigned> Metadata;
unsigned Partition = NoPartition;
};
struct MetadataNode {
StringRef Line;
std::vector<unsigned> Refs;
std::vector<unsigned> Symbols;
};
bool isNameChar(char C) {
return (C >= 'a' && C <= 'z') || (C >= 'A' && C <= 'Z') ||
(C >= '0' && C <= '9') || C == '-' || C == '$' || C == '.' ||
C == '_';
}
bool isDigit(char C) { return C >= '0' && C <= '9'; }
// Length of the unquoted symbol name starting at Text[Start], or 0.
size_t nameLength(StringRef Text, size_t Start) {
size_t End = Start;
while (End < Text.size() && isNameChar(Text[End]))
++End;
return End - Start;
}
uint32_t stableHash(StringRef Name) {
uint32_t Hash = 0x811c9dc5u;
for (char C : Name) {
Hash ^= static_cast<unsigned char>(C);
Hash *= 0x01000193u;
}
return Hash;
}
bool startsWithLinkage(StringRef Rest) {
for (StringRef Linkage :
{"linkonce", "weak", "common", "appending", "extern_weak",
"available_externally", "external"})
if (Rest.starts_with(Linkage))
return true;
return false;
}
bool hasVisibilityToken(StringRef Prefix) {
SmallVector<StringRef, 8> Tokens;
Prefix.split(Tokens, ' ', -1, false);
for (StringRef Token : Tokens)
if (Token == "dso_local" || Token == "dso_preemptable" ||
Token == "hidden" || Token == "protected" || Token == "default" ||
Token == "dllimport" || Token == "dllexport")
return true;
return false;
}
// Position of the first unquoted occurrence of Wanted, or npos.
size_t findUnquoted(StringRef Text, char Wanted) {
bool Quoted = false;
for (size_t I = 0; I < Text.size(); ++I) {
if (Text[I] == '"')
Quoted = !Quoted;
else if (!Quoted && Text[I] == Wanted)
return I;
}
return StringRef::npos;
}
bool parseFunctionHeader(StringRef Line, Symbol &S) {
StringRef Rest = Line.drop_front(sizeof("define ") - 1);
if (Rest.consume_front("internal ") || Rest.consume_front("private "))
S.Local = true;
else if (startsWithLinkage(Rest))
return false;
size_t At = Rest.find('@');
if (At == StringRef::npos)
return false;
StringRef Prefix = Rest.take_front(At);
if (S.Local && hasVisibilityToken(Prefix))
return false;
size_t Length = nameLength(Rest, At + 1);
size_t Open = At + 1 + Length;
if (Length == 0 || Open >= Rest.size() || Rest[Open] != '(')
return false;
S.Name = Rest.substr(At + 1, Length);
size_t Depth = 0, Close = StringRef::npos;
bool Quoted = false;
for (size_t I = Open; I < Rest.size(); ++I) {
char C = Rest[I];
if (C == '"')
Quoted = !Quoted;
if (Quoted)
continue;
if (C == '(')
++Depth;
else if (C == ')' && --Depth == 0) {
Close = I;
break;
}
}
if (Close == StringRef::npos)
return false;
StringRef Tail = Rest.drop_front(Close + 1);
size_t Brace = findUnquoted(Tail, '{');
if (Brace == StringRef::npos)
return false;
StringRef After = Tail.drop_front(Brace + 1).ltrim();
if (!After.empty() && After.front() != ';')
return false;
S.Declaration = "declare ";
if (S.Local)
S.Declaration += "hidden ";
S.Declaration += Prefix;
S.Declaration += '@';
S.Declaration += S.Name;
S.Declaration += Rest.slice(Open, Close + 1);
// Attribute groups and address significance carry over. Personality,
// sections, comdats and debug attachments belong only to the definition.
SmallVector<StringRef, 8> Tokens;
Tail.take_front(Brace).split(Tokens, ' ', -1, false);
for (StringRef Token : Tokens) {
bool Group = Token.size() > 1 && Token.front() == '#' &&
all_of(Token.drop_front(), isDigit);
if (Group || Token == "unnamed_addr" || Token == "local_unnamed_addr") {
S.Declaration += ' ';
S.Declaration += Token;
}
}
if (S.Local)
S.PromotedHeader = (Twine("define hidden ") + Rest).str();
return true;
}
// Parses `@name = ...`. Returns false for unsupported forms.
bool parseGlobal(StringRef Line, Symbol &S) {
size_t Length = nameLength(Line, 1);
if (Length == 0)
return false;
S.Name = Line.substr(1, Length);
StringRef Rest = Line.drop_front(1 + Length);
if (!Rest.consume_front(" = "))
return false;
if (Rest.starts_with("external ") || Rest.starts_with("extern_weak ")) {
S.Kind = SymbolKind::Declaration;
return true;
}
S.Kind = SymbolKind::Variable;
if (Rest.consume_front("internal ") || Rest.consume_front("private "))
S.Local = true;
else if (startsWithLinkage(Rest))
return false;
StringRef Cursor = Rest;
std::string Qualifiers;
StringRef Keyword;
while (true) {
if (Cursor.consume_front("global ")) {
Keyword = "global";
break;
}
if (Cursor.consume_front("constant ")) {
Keyword = "constant";
break;
}
size_t Space = Cursor.find(' ');
if (Space == StringRef::npos)
return false;
StringRef Token = Cursor.take_front(Space);
if (Token == "dso_local" || Token == "hidden" || Token == "protected" ||
Token == "default") {
if (S.Local)
return false;
Qualifiers += Token;
Qualifiers += ' ';
} else if (Token.starts_with("thread_local") ||
Token.starts_with("addrspace(")) {
Qualifiers += Token;
Qualifiers += ' ';
} else if (Token != "unnamed_addr" && Token != "local_unnamed_addr" &&
Token != "externally_initialized") {
return false;
}
Cursor = Cursor.drop_front(Space + 1);
}
// The value type ends at the first space outside brackets.
int Depth = 0;
size_t TypeEnd = StringRef::npos;
for (size_t I = 0; I < Cursor.size(); ++I) {
char C = Cursor[I];
if (C == '{' || C == '[' || C == '<' || C == '(')
++Depth;
else if (C == '}' || C == ']' || C == '>' || C == ')')
--Depth;
else if (C == ' ' && Depth == 0) {
TypeEnd = I;
break;
}
}
if (TypeEnd == StringRef::npos)
return false;
S.Declaration = (Twine("@") + S.Name + " = external " + (S.Local ? "hidden " : "") +
Qualifiers + Keyword + " " + Cursor.take_front(TypeEnd))
.str();
if (S.Local)
S.PromotedHeader = (Twine("@") + S.Name + " = hidden " + Rest).str();
return true;
}
// Collects the symbols (sorted) and metadata nodes (in order of first
// appearance) that Text references outside quoted strings.
void scanReferences(StringRef Text, const StringMap<unsigned> &Index,
std::vector<unsigned> &Refs,
std::vector<unsigned> &Metadata) {
DenseSet<unsigned> Seen(Metadata.begin(), Metadata.end());
bool Quoted = false;
for (size_t I = 0; I < Text.size(); ++I) {
char C = Text[I];
if (C == '\n') {
Quoted = false;
} else if (C == '"') {
Quoted = !Quoted;
} else if (Quoted) {
continue;
} else if (C == '@') {
size_t Length = nameLength(Text, I + 1);
if (Length != 0) {
auto It = Index.find(Text.substr(I + 1, Length));
if (It != Index.end())
Refs.push_back(It->second);
I += Length;
}
} else if (C == '!' && I + 1 < Text.size() && isDigit(Text[I + 1])) {
unsigned Number = 0;
size_t J = I + 1;
for (; J < Text.size() && isDigit(Text[J]); ++J)
Number = Number * 10 + static_cast<unsigned>(Text[J] - '0');
if (Seen.insert(Number).second)
Metadata.push_back(Number);
I = J - 1;
}
}
llvm::sort(Refs);
Refs.erase(std::unique(Refs.begin(), Refs.end()), Refs.end());
}
// Appends Text with every unquoted metadata reference renumbered by Map.
void appendRenumbered(StringRef Text, const DenseMap<unsigned, unsigned> &Map,
std::string &Out) {
bool Quoted = false;
size_t Copied = 0;
for (size_t I = 0; I < Text.size(); ++I) {
char C = Text[I];
if (C == '\n') {
Quoted = false;
} else if (C == '"') {
Quoted = !Quoted;
} else if (!Quoted && C == '!' && I + 1 < Text.size() &&
isDigit(Text[I + 1])) {
unsigned Number = 0;
size_t J = I + 1;
for (; J < Text.size() && isDigit(Text[J]); ++J)
Number = Number * 10 + static_cast<unsigned>(Text[J] - '0');
auto It = Map.find(Number);
if (It != Map.end()) {
Out.append(Text.data() + Copied, I + 1 - Copied);
Out += std::to_string(It->second);
Copied = J;
}
I = J - 1;
}
}
Out.append(Text.data() + Copied, Text.size() - Copied);
}
} // namespace
std::optional<ProgramPartitions> partitionProgramText(StringRef Text,
unsigned Count) {
if (Count < 2)
return std::nullopt;
std::vector<StringRef> Header, Types, Attributes, NamedMetadata;
std::vector<unsigned> CompileUnits;
DenseMap<unsigned, MetadataNode> Metadata;
std::vector<Symbol> Symbols;
StringMap<unsigned> Index;
auto Add = [&](Symbol S) {
if (!Index.try_emplace(S.Name, Symbols.size()).second)
return false;
Symbols.push_back(std::move(S));
return true;
};
size_t Pos = 0;
while (Pos < Text.size()) {
size_t End = Text.find('\n', Pos);
if (End == StringRef::npos)
End = Text.size();
StringRef Line = Text.slice(Pos, End);
size_t Next = End + 1;
if (Line.empty() || Line.trim().empty()) {
// Blank separator.
} else if (Line.starts_with("define ")) {
Symbol S;
S.Kind = SymbolKind::Function;
S.Header = Line;
if (!parseFunctionHeader(Line, S))
return std::nullopt;
// The emitter indents every body line, so the first line that is
// exactly "}" closes the function.
size_t Close = End;
while (true) {
Close = Text.find("\n}", Close);
if (Close == StringRef::npos)
return std::nullopt;
if (Close + 2 == Text.size() || Text[Close + 2] == '\n')
break;
Close += 2;
}
S.Body = Text.slice(std::min(End + 1, Text.size()), Close + 2);
Next = Close + 3;
if (!Add(std::move(S)))
return std::nullopt;
} else if (Line.front() == '@') {
Symbol S;
S.Header = Line;
if (!parseGlobal(Line, S) || !Add(std::move(S)))
return std::nullopt;
} else if (Line.starts_with("declare ")) {
Symbol S;
S.Header = Line;
size_t At = Line.find('@');
size_t Length = At == StringRef::npos ? 0 : nameLength(Line, At + 1);
if (Length == 0)
return std::nullopt;
S.Name = Line.substr(At + 1, Length);
if (!Add(std::move(S)))
return std::nullopt;
} else if (Line.front() == '%') {
if (Line.find(" = type ") == StringRef::npos)
return std::nullopt;
Types.push_back(Line);
} else if (Line.starts_with("attributes #")) {
Attributes.push_back(Line);
} else if (Line.front() == '!') {
if (Line.size() > 1 && isDigit(Line[1])) {
unsigned Number = 0;
size_t I = 1;
for (; I < Line.size() && isDigit(Line[I]); ++I)
Number = Number * 10 + static_cast<unsigned>(Line[I] - '0');
if (!Line.drop_front(I).starts_with(" = "))
return std::nullopt;
if (!Metadata.try_emplace(Number, MetadataNode{Line, {}, {}}).second)
return std::nullopt;
} else if (Line.starts_with("!llvm.dbg.cu = !{")) {
// Each partition lists only the compile units its functions use.
std::vector<unsigned> Unused;
scanReferences(Line, Index, Unused, CompileUnits);
} else {
NamedMetadata.push_back(Line);
}
} else if (Line.front() == ';' || Line.starts_with("source_filename ") ||
Line.starts_with("target ")) {
Header.push_back(Line);
} else {
return std::nullopt;
}
Pos = Next;
}
for (Symbol &S : Symbols) {
if (S.Kind == SymbolKind::Declaration)
continue;
scanReferences(S.Header, Index, S.Refs, S.Metadata);
if (!S.Body.empty()) {
std::vector<unsigned> Refs;
scanReferences(S.Body, Index, Refs, S.Metadata);
S.Refs.insert(S.Refs.end(), Refs.begin(), Refs.end());
llvm::sort(S.Refs);
S.Refs.erase(std::unique(S.Refs.begin(), S.Refs.end()), S.Refs.end());
}
}
// An internal function referenced by exactly one other function, and by
// no global, is a whole-program inlining candidate; keep it with that
// function so its partition can inline it without importing. A function
// body larger than this is past LLVM's inline cost limit even for the only
// call to a local function, so it is placed by its own name. Otherwise a
// program's module initializers, which only the entry point calls, would
// all share the entry point's partition.
constexpr size_t MaxInlinedBodyBytes = 256 * 1024;
constexpr unsigned Shared = ~1u;
std::vector<unsigned> Referrer(Symbols.size(), NoPartition);
for (unsigned I = 0; I < Symbols.size(); ++I) {
const Symbol &S = Symbols[I];
if (S.Kind == SymbolKind::Declaration)
continue;
for (unsigned Ref : S.Refs) {
if (Ref == I)
continue;
if (S.Kind == SymbolKind::Variable || Referrer[Ref] != NoPartition)
Referrer[Ref] = Shared;
else
Referrer[Ref] = I;
}
}
for (unsigned I = 0; I < Symbols.size(); ++I) {
if (Symbols[I].Kind != SymbolKind::Function)
continue;
unsigned Root = I;
for (unsigned Steps = 0; Steps < 1024; ++Steps) {
const Symbol &S = Symbols[Root];
unsigned Parent = Referrer[Root];
if (!S.Local || S.Body.size() > MaxInlinedBodyBytes ||
Parent == NoPartition || Parent == Shared ||
Symbols[Parent].Kind != SymbolKind::Function || Parent == I)
break;
Root = Parent;
}
Symbols[I].Partition = stableHash(Symbols[Root].Name) % Count;
}
for (const Symbol &S : Symbols) {
if (S.Kind != SymbolKind::Function)
continue;
for (unsigned Ref : S.Refs)
if (Symbols[Ref].Kind == SymbolKind::Variable &&
Symbols[Ref].Partition == NoPartition)
Symbols[Ref].Partition = S.Partition;
}
for (Symbol &S : Symbols)
if (S.Kind == SymbolKind::Variable && S.Partition == NoPartition)
S.Partition = stableHash(S.Name) % Count;
ProgramPartitions Result;
for (const Symbol &S : Symbols)
if (S.Kind != SymbolKind::Declaration && !S.Local)
Result.PublicSymbols.insert(S.Name);
// Metadata nodes are numbered in the emitter's global order, so one added
// node would renumber the references in every partition. Each partition
// instead numbers the nodes it uses in order of first use, which keeps an
// unchanged partition's text, and therefore its cached artifacts, intact.
for (auto &Entry : Metadata) {
MetadataNode &Node = Entry.second;
size_t Equals = Node.Line.find(" = ");
scanReferences(Node.Line.drop_front(Equals + 3), Index, Node.Symbols,
Node.Refs);
}
std::vector<unsigned> NamedRoots, NamedSymbols;
for (StringRef Line : NamedMetadata)
scanReferences(Line, Index, NamedSymbols, NamedRoots);
Result.Sources.resize(Count);
std::vector<std::vector<unsigned>> Owned(Count);
for (unsigned I = 0; I < Symbols.size(); ++I)
if (Symbols[I].Kind != SymbolKind::Declaration)
Owned[Symbols[I].Partition].push_back(I);
auto Assemble = [&](unsigned P) {
std::vector<char> Needed(Symbols.size(), 0);
DenseMap<unsigned, unsigned> Renumber;
std::vector<unsigned> Nodes;
auto Use = [&](ArrayRef<unsigned> Numbers) {
for (unsigned Number : Numbers)
if (Metadata.count(Number) &&
Renumber.try_emplace(Number, Nodes.size()).second)
Nodes.push_back(Number);
};
Use(NamedRoots);
for (unsigned Ref : NamedSymbols)
Needed[Ref] = 1;
size_t Bytes = 0;
for (SymbolKind Kind : {SymbolKind::Variable, SymbolKind::Function})
for (unsigned I : Owned[P])
if (Symbols[I].Kind == Kind) {
const Symbol &S = Symbols[I];
for (unsigned Ref : S.Refs)
Needed[Ref] = 1;
Use(S.Metadata);
Bytes += S.Header.size() + S.Body.size() + 2;
}
for (size_t K = 0; K < Nodes.size(); ++K) {
const MetadataNode &Node = Metadata.find(Nodes[K])->second;
Use(Node.Refs);
for (unsigned Ref : Node.Symbols)
Needed[Ref] = 1;
}
std::string &Out = Result.Sources[P];
Out.reserve(Bytes + (Bytes >> 3));
auto Line = [&](StringRef Text) {
if (Renumber.empty())
Out += Text;
else
appendRenumbered(Text, Renumber, Out);
Out += '\n';
};
for (StringRef Text : Header)
Line(Text);
for (StringRef Text : Types)
Line(Text);
// Declarations follow name order rather than the emitter's first-use
// order, which an edit elsewhere in the program can change.
std::vector<unsigned> Declared;
for (unsigned I = 0; I < Symbols.size(); ++I)
if (Needed[I] && (Symbols[I].Kind == SymbolKind::Declaration ||
Symbols[I].Partition != P))
Declared.push_back(I);
llvm::sort(Declared, [&](unsigned A, unsigned B) {
return Symbols[A].Name < Symbols[B].Name;
});
for (unsigned I : Declared) {
const Symbol &S = Symbols[I];
Line(S.Kind == SymbolKind::Declaration ? S.Header
: StringRef(S.Declaration));
}
for (SymbolKind Kind : {SymbolKind::Variable, SymbolKind::Function})
for (unsigned I : Owned[P]) {
const Symbol &S = Symbols[I];
if (S.Kind != Kind)
continue;
Line(S.Local ? StringRef(S.PromotedHeader) : S.Header);
if (!S.Body.empty())
Line(S.Body);
}
for (StringRef Text : Attributes)
Line(Text);
for (StringRef Text : NamedMetadata)
Line(Text);
std::string Units;
for (unsigned Unit : CompileUnits) {
auto It = Renumber.find(Unit);
if (It == Renumber.end())
continue;
Units += Units.empty() ? "!" : ", !";
Units += std::to_string(It->second);
}
if (!Units.empty())
Out += "!llvm.dbg.cu = !{" + Units + "}\n";
for (size_t K = 0; K < Nodes.size(); ++K) {
StringRef Text = Metadata.find(Nodes[K])->second.Line;
Out += '!';
Out += std::to_string(K);
Line(Text.drop_front(Text.find(" = ")));
}
};
DefaultThreadPool Pool(hardware_concurrency());
for (unsigned P = 0; P < Count; ++P)
Pool.async([&Assemble, P] { Assemble(P); });
Pool.wait();
return Result;
}
} // namespace scriptc