Files
scriptc/native/llvm-codegen/src/partition.cpp
T
Chris Tate 3ea231d70d Optimize and generate native code in parallel partitions
Large executables were simplified and inlined as one module before the
optimizer split them, so most of a release build ran on a single core,
and any edit recompiled the whole program.

The LLVM helper now divides the program before optimization. It splits
the emitter's assembly text directly, so partitions also parse
concurrently, and falls back to LLVM's module splitter for any input
outside the emitter's grammar. Internal functions referenced by a single
function stay with that function, and globals stay with their first user.
Optimized partitions follow the ThinLTO model: each is simplified and
summarized on its own thread, a thin link imports small functions across
partitions for inlining and returns unexported hidden definitions to
internal linkage, then each partition is optimized and compiled on its
own thread. Unoptimized partitions compile independently, so dev builds
of large programs are parallel too.

A private per-helper cache directory stores simplified partitions under
the hash of their source and objects under the hash of the module after
importing, so a rebuild recompiles only the partitions an edit can
affect. To keep unchanged partitions identical across edits, each
partition numbers its metadata in first-use order, lists only the debug
compile units it uses, and orders declarations by name, and interned
string constants are named by their content rather than first-use order.

Partition counts are powers of two that depend only on the module size,
up to 64, so output stays the same on every host and functions keep
their partition across ordinary edits. LLVM text that arrives in parts is
written in large batches and handed to the helper in place instead of
being written a second time.
2026-10-09 19:51:47 -05:00

584 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.
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 || 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