Skip to content

Commit

Permalink
Always call MemUsage::Collect to collect metrics from a field (#4480)
Browse files Browse the repository at this point in the history
Previously Collect() was used for types that implemented
CollectMemUsage() but otherwise Add() was used. This required the caller
to think about the type of the field and know/decide which method to
use.

Now, the caller always uses Collect() unless they are adding specific
byte values, in which case Add is used. Typically then, Add will only be
used to implement the CollectMemUsage() function.

To do this we require all Collect() methods to be templates so that they
all be a single overload set. The Collect on BumpPtrAllocator is
converted to a template that checks
`std::same_as<llvm::BumpPtrAllocator, T>`.
  • Loading branch information
danakj authored Nov 5, 2024
1 parent f03bd6a commit 361efa9
Show file tree
Hide file tree
Showing 13 changed files with 42 additions and 35 deletions.
26 changes: 16 additions & 10 deletions toolchain/base/mem_usage.h
Original file line number Diff line number Diff line change
Expand Up @@ -47,27 +47,28 @@ class MemUsage {
.reserved_bytes = reserved_bytes});
}

// Adds usage tracking for an allocator.
auto Add(std::string label, const llvm::BumpPtrAllocator& allocator) -> void {
// Adds memory usage for a `llvm::BumpPtrAllocator`.
auto Collect(std::string label, const llvm::BumpPtrAllocator& allocator)
-> void {
Add(std::move(label), allocator.getBytesAllocated(),
allocator.getTotalMemory());
}

// Adds usage tracking for a map.
// Adds memory usage for a `Map`.
template <typename KeyT, typename ValueT, ssize_t SmallSize,
typename KeyContextT>
auto Add(std::string label, Map<KeyT, ValueT, SmallSize, KeyContextT> map,
KeyContextT key_context = KeyContextT()) -> void {
auto Collect(std::string label, Map<KeyT, ValueT, SmallSize, KeyContextT> map,
KeyContextT key_context = KeyContextT()) -> void {
// These don't track used bytes, so we set the same value for used and
// reserved bytes.
auto bytes = map.ComputeMetrics(key_context).storage_bytes;
Add(std::move(label), bytes, bytes);
}

// Adds usage tracking for a set.
// Adds memory usage for a `Set`.
template <typename KeyT, ssize_t SmallSize, typename KeyContextT>
auto Add(std::string label, Set<KeyT, SmallSize, KeyContextT> set,
KeyContextT key_context = KeyContextT()) -> void {
auto Collect(std::string label, Set<KeyT, SmallSize, KeyContextT> set,
KeyContextT key_context = KeyContextT()) -> void {
// These don't track used bytes, so we set the same value for used and
// reserved bytes.
auto bytes = set.ComputeMetrics(key_context).storage_bytes;
Expand All @@ -81,7 +82,8 @@ class MemUsage {
// This uses SmallVector in order to get proper inference for T, which
// ArrayRef misses.
template <typename T, unsigned N>
auto Add(std::string label, const llvm::SmallVector<T, N>& array) -> void {
auto Collect(std::string label, const llvm::SmallVector<T, N>& array)
-> void {
Add(std::move(label), array.size_in_bytes(), array.capacity_in_bytes());
}

Expand All @@ -90,7 +92,11 @@ class MemUsage {
// The expected signature of `CollectMemUsage` is above, in MemUsage class
// comments.
template <typename T>
auto Collect(llvm::StringRef label, const T& arg) -> void {
requires requires(MemUsage& mem_usage, llvm::StringRef label,
const T& arg) {
{ arg.CollectMemUsage(mem_usage, label) };
}
auto Collect(std::string label, const T& arg) -> void {
arg.CollectMemUsage(*this, label);
}

Expand Down
2 changes: 1 addition & 1 deletion toolchain/base/value_store.h
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ class ValueStore
// Collects memory usage of the values.
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(label.str(), values_);
mem_usage.Collect(label.str(), values_);
}

auto array_ref() const -> llvm::ArrayRef<ValueType> { return values_; }
Expand Down
8 changes: 4 additions & 4 deletions toolchain/lex/tokenized_buffer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -373,10 +373,10 @@ auto TokenizedBuffer::AddComment(int32_t indent, int32_t start, int32_t end)

auto TokenizedBuffer::CollectMemUsage(MemUsage& mem_usage,
llvm::StringRef label) const -> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
mem_usage.Add(MemUsage::ConcatLabel(label, "token_infos_"), token_infos_);
mem_usage.Add(MemUsage::ConcatLabel(label, "line_infos_"), line_infos_);
mem_usage.Add(MemUsage::ConcatLabel(label, "comments_"), comments_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "token_infos_"), token_infos_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "line_infos_"), line_infos_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "comments_"), comments_);
}

auto TokenizedBuffer::SourceBufferDiagnosticConverter::ConvertLoc(
Expand Down
4 changes: 2 additions & 2 deletions toolchain/parse/tree.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -71,8 +71,8 @@ auto Tree::Verify() const -> ErrorOr<Success> {

auto Tree::CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "node_impls_"), node_impls_);
mem_usage.Add(MemUsage::ConcatLabel(label, "imports_"), imports_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "node_impls_"), node_impls_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "imports_"), imports_);
}

auto Tree::PostorderIterator::MakeRange(NodeId begin, NodeId end)
Expand Down
3 changes: 2 additions & 1 deletion toolchain/parse/tree_and_subtrees.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -235,7 +235,8 @@ auto TreeAndSubtrees::PrintPreorder(llvm::raw_ostream& output) const -> void {

auto TreeAndSubtrees::CollectMemUsage(MemUsage& mem_usage,
llvm::StringRef label) const -> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "subtree_sizes_"), subtree_sizes_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "subtree_sizes_"),
subtree_sizes_);
}

auto TreeAndSubtrees::SiblingIterator::Print(llvm::raw_ostream& output) const
Expand Down
4 changes: 2 additions & 2 deletions toolchain/sem_ir/block_value_store.h
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,8 @@ class BlockValueStore : public Yaml::Printable<BlockValueStore<IdT>> {
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
mem_usage.Add(MemUsage::ConcatLabel(label, "canonical_blocks_"),
canonical_blocks_, KeyContext(this));
mem_usage.Collect(MemUsage::ConcatLabel(label, "canonical_blocks_"),
canonical_blocks_, KeyContext(this));
}

auto size() const -> int { return values_.size(); }
Expand Down
10 changes: 5 additions & 5 deletions toolchain/sem_ir/constant.h
Original file line number Diff line number Diff line change
Expand Up @@ -105,9 +105,9 @@ class ConstantValueStore {
// Collects memory usage of members.
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "values_"), values_);
mem_usage.Add(MemUsage::ConcatLabel(label, "symbolic_constants_"),
symbolic_constants_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "symbolic_constants_"),
symbolic_constants_);
}

// Returns the constant values mapping as an ArrayRef whose keys are
Expand Down Expand Up @@ -155,8 +155,8 @@ class ConstantStore {
// Collects memory usage of members.
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "map_"), map_);
mem_usage.Add(MemUsage::ConcatLabel(label, "constants_"), constants_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "map_"), map_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "constants_"), constants_);
}

// Returns a copy of the constant IDs as a vector, in an arbitrary but
Expand Down
2 changes: 1 addition & 1 deletion toolchain/sem_ir/file.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ auto File::OutputYaml(bool include_builtins) const -> Yaml::OutputMapping {

auto File::CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "allocator_"), allocator_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "entity_names_"),
entity_names_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "functions_"), functions_);
Expand Down
4 changes: 2 additions & 2 deletions toolchain/sem_ir/generic.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ auto SpecificStore::GetOrAdd(GenericId generic_id, InstBlockId args_id)
auto SpecificStore::CollectMemUsage(MemUsage& mem_usage,
llvm::StringRef label) const -> void {
mem_usage.Collect(MemUsage::ConcatLabel(label, "specifics_"), specifics_);
mem_usage.Add(MemUsage::ConcatLabel(label, "lookup_table_"), lookup_table_,
KeyContext(specifics_.array_ref()));
mem_usage.Collect(MemUsage::ConcatLabel(label, "lookup_table_"),
lookup_table_, KeyContext(specifics_.array_ref()));
}

auto GetConstantInSpecific(const File& sem_ir, SpecificId specific_id,
Expand Down
2 changes: 1 addition & 1 deletion toolchain/sem_ir/impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ class ImplStore {
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
mem_usage.Add(MemUsage::ConcatLabel(label, "lookup_"), lookup_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "lookup_"), lookup_);
}

auto array_ref() const -> llvm::ArrayRef<Impl> { return values_.array_ref(); }
Expand Down
2 changes: 1 addition & 1 deletion toolchain/sem_ir/inst.h
Original file line number Diff line number Diff line change
Expand Up @@ -423,7 +423,7 @@ class InstStore {
// Collects memory usage of members.
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "loc_ids_"), loc_ids_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "loc_ids_"), loc_ids_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
}

Expand Down
2 changes: 1 addition & 1 deletion toolchain/sem_ir/name_scope.h
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ class NameScopeStore {
// Collects memory usage of members.
auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Collect(label, values_);
mem_usage.Collect(std::string(label), values_);
}

private:
Expand Down
8 changes: 4 additions & 4 deletions toolchain/sem_ir/type.h
Original file line number Diff line number Diff line change
Expand Up @@ -126,10 +126,10 @@ class TypeStore : public Yaml::Printable<TypeStore> {

auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
-> void {
mem_usage.Add(MemUsage::ConcatLabel(label, "complete_type_info_"),
complete_type_info_);
mem_usage.Add(MemUsage::ConcatLabel(label, "complete_types_"),
complete_types_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "complete_type_info_"),
complete_type_info_);
mem_usage.Collect(MemUsage::ConcatLabel(label, "complete_types_"),
complete_types_);
}

private:
Expand Down

0 comments on commit 361efa9

Please sign in to comment.