Add basic output of where memory is stored after a compile. (#4136)

The output is really basic, I'm just adding this to help track how
memory is allocated.

```
---
filename:        'check/testdata/expr_category/in_place_tuple_init.carbon'
source_:
  used_bytes:      8057
  reserved_bytes:  8057
tokens_.allocator_:
  used_bytes:      0
  reserved_bytes:  0
tokens_.token_infos_:
  used_bytes:      1040
  reserved_bytes:  2032

(eliding)

value_stores_.string_literals_.set_:
  used_bytes:      320
  reserved_bytes:  320
Total:
  used_bytes:      20609
  reserved_bytes:  29437
...
```

---------

Co-authored-by: josh11b <15258583+josh11b@users.noreply.github.com>

toolchain/base/BUILD

@@ -31,11 +31,23 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "mem_usage",
+    hdrs = ["mem_usage.h"],
+    deps = [
+        ":yaml",
+        "//common:map",
+        "//common:set",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "value_store",
     hdrs = ["value_store.h"],
     deps = [
         ":index_base",
+        ":mem_usage",
         ":yaml",
         "//common:check",
         "//common:hashing",

toolchain/base/mem_usage.h

@@ -0,0 +1,145 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef CARBON_TOOLCHAIN_BASE_MEM_USAGE_H_
+#define CARBON_TOOLCHAIN_BASE_MEM_USAGE_H_
+
+#include <cstdint>
+
+#include "common/map.h"
+#include "common/set.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "toolchain/base/yaml.h"
+
+namespace Carbon {
+
+// Helps track memory usage for a compile.
+//
+// Users will mix `Add` and `Collect` calls, using `ConcatLabel` to label
+// allocation sources. Typically we'll collect stats for growable, potentially
+// large data types (such as `SmallVector`), ignoring small fixed-size members
+// (such as pointers or `int32_t`).
+//
+// For example:
+//
+//   auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+//       -> void {
+//     // Explicit tracking.
+//     mem_usage.Add(MemUsage::ConcatLabel(label, "data_"), data_.used_bytes(),
+//                   data_.reserved_bytes());
+//     // Common library types like `Map` and `llvm::SmallVector` have
+//     // type-specific support.
+//     mem_usage.Add(MemUsage::Concat(label, "array_"), array_);
+//     // Implementing `CollectMemUsage` allows use with the same interface.
+//     mem_usage.Collect(MemUsage::Concat(label, "obj_"), obj_);
+//   }
+class MemUsage {
+ public:
+  // Adds tracking for used and reserved bytes, paired with the given label.
+  auto Add(std::string label, int64_t used_bytes, int64_t reserved_bytes)
+      -> void {
+    mem_usage_.push_back({.label = std::move(label),
+                          .used_bytes = used_bytes,
+                          .reserved_bytes = reserved_bytes});
+  }
+
+  // Adds usage tracking for an allocator.
+  auto Add(std::string label, const llvm::BumpPtrAllocator& allocator) -> void {
+    Add(std::move(label), allocator.getBytesAllocated(),
+        allocator.getTotalMemory());
+  }
+
+  // Adds usage tracking 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 {
+    // 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.
+  template <typename KeyT, ssize_t SmallSize, typename KeyContextT>
+  auto Add(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;
+    Add(std::move(label), bytes, bytes);
+  }
+
+  // Adds memory usage of an array's data. This ignores the possible overhead of
+  // a SmallVector's in-place storage; if it's used, it's going to be tiny
+  // relative to scaling memory costs.
+  //
+  // 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 {
+    Add(std::move(label), array.size_in_bytes(), array.capacity_in_bytes());
+  }
+
+  // Adds memory usage for an object that provides `CollectMemUsage`.
+  //
+  // The expected signature of `CollectMemUsage` is above, in MemUsage class
+  // comments.
+  template <typename T>
+  auto Collect(llvm::StringRef label, const T& arg) -> void {
+    arg.CollectMemUsage(*this, label);
+  }
+
+  // Constructs a label for memory usage, handling the `.` concatenation.
+  // We don't expect much depth in labels per-call.
+  static auto ConcatLabel(llvm::StringRef label, llvm::StringRef child_label)
+      -> std::string {
+    return llvm::formatv("{0}.{1}", label, child_label);
+  }
+  static auto ConcatLabel(llvm::StringRef label, llvm::StringRef child_label1,
+                          llvm::StringRef child_label2) -> std::string {
+    return llvm::formatv("{0}.{1}.{2}", label, child_label1, child_label2);
+  }
+
+  auto OutputYaml(llvm::StringRef filename) const -> Yaml::OutputMapping {
+    // Explicitly copy the filename.
+    return Yaml::OutputMapping([&, filename](Yaml::OutputMapping::Map map) {
+      map.Add("filename", filename);
+      int64_t total_used = 0;
+      int64_t total_reserved = 0;
+      for (const auto& entry : mem_usage_) {
+        total_used += entry.used_bytes;
+        total_reserved += entry.reserved_bytes;
+        map.Add(entry.label,
+                Yaml::OutputMapping([&](Yaml::OutputMapping::Map byte_map) {
+                  byte_map.Add("used_bytes", entry.used_bytes);
+                  byte_map.Add("reserved_bytes", entry.reserved_bytes);
+                }));
+      }
+      map.Add("Total",
+              Yaml::OutputMapping([&](Yaml::OutputMapping::Map byte_map) {
+                byte_map.Add("used_bytes", total_used);
+                byte_map.Add("reserved_bytes", total_reserved);
+              }));
+    });
+  }
+
+ private:
+  // Memory usage for a specific label.
+  struct Entry {
+    std::string label;
+    int64_t used_bytes;
+    int64_t reserved_bytes;
+  };
+
+  // The accumulated data on memory usage.
+  llvm::SmallVector<Entry> mem_usage_;
+};
+
+}  // namespace Carbon
+
+#endif  // CARBON_TOOLCHAIN_BASE_MEM_USAGE_H_

toolchain/base/value_store.h

@@ -17,6 +17,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/YAMLParser.h"
 #include "toolchain/base/index_base.h"
+#include "toolchain/base/mem_usage.h"
 #include "toolchain/base/yaml.h"
 
 namespace Carbon {
@@ -187,6 +188,12 @@ class ValueStore
     });
   }
 
+  // Collects memory usage of the values.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void {
+    mem_usage.Add(label.str(), values_);
+  }
+
   auto array_ref() const -> llvm::ArrayRef<ValueType> { return values_; }
   auto size() const -> size_t { return values_.size(); }
 
@@ -237,6 +244,15 @@ class CanonicalValueStore {
   }
   auto size() const -> size_t { return values_.size(); }
 
+  // Collects memory usage of the values and deduplication set.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void {
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
+    auto bytes =
+        set_.ComputeMetrics(KeyContext(values_.array_ref())).storage_bytes;
+    mem_usage.Add(MemUsage::ConcatLabel(label, "set_"), bytes, bytes);
+  }
+
  private:
   class KeyContext;
 
@@ -322,6 +338,18 @@ class SharedValueStores : public Yaml::Printable<SharedValueStores> {
     });
   }
 
+  // Collects memory usage for the various shared stores.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void {
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "ints_"), ints_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "reals_"), reals_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "floats_"), floats_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "identifiers_"),
+                      identifiers_);
+    mem_usage.Collect(MemUsage::ConcatLabel(label, "string_literals_"),
+                      string_literals_);
+  }
+
  private:
   CanonicalValueStore<IntId> ints_;
   ValueStore<RealId> reals_;

toolchain/driver/driver.cpp

@@ -308,6 +308,14 @@ Dump the generated assembly to stdout after codegen.
 )""",
         },
         [&](auto& arg_b) { arg_b.Set(&dump_asm); });
+    b.AddFlag(
+        {
+            .name = "dump-mem-usage",
+            .help = R"""(
+Dumps the amount of memory used.
+)""",
+        },
+        [&](auto& arg_b) { arg_b.Set(&dump_mem_usage); });
     b.AddFlag(
         {
             .name = "prelude-import",
@@ -344,6 +352,7 @@ Excludes files with the given prefix from dumps.
   bool dump_sem_ir = false;
   bool dump_llvm_ir = false;
   bool dump_asm = false;
+  bool dump_mem_usage = false;
   bool stream_errors = false;
   bool preorder_parse_tree = false;
   bool builtin_sem_ir = false;
@@ -540,6 +549,9 @@ class Driver::CompilationUnit {
       sorting_consumer_ = SortingDiagnosticConsumer(*consumer);
       consumer_ = &*sorting_consumer_;
     }
+    if (options_.dump_mem_usage && IncludeInDumps()) {
+      mem_usage_ = MemUsage();
+    }
   }
 
   // Loads source and lexes it. Returns true on success.
@@ -552,6 +564,10 @@ class Driver::CompilationUnit {
                                              *consumer_);
       }
     });
+    if (mem_usage_) {
+      mem_usage_->Add("source_", source_->text().size(),
+                      source_->text().size());
+    }
     if (!source_) {
       success_ = false;
       return;
@@ -565,6 +581,9 @@ class Driver::CompilationUnit {
       consumer_->Flush();
       driver_->output_stream_ << tokens_;
     }
+    if (mem_usage_) {
+      mem_usage_->Collect("tokens_", *tokens_);
+    }
     CARBON_VLOG() << "*** Lex::TokenizedBuffer ***\n" << tokens_;
     if (tokens_->has_errors()) {
       success_ = false;
@@ -582,6 +601,9 @@ class Driver::CompilationUnit {
       consumer_->Flush();
       parse_tree_->Print(driver_->output_stream_, options_.preorder_parse_tree);
     }
+    if (mem_usage_) {
+      mem_usage_->Collect("parse_tree_", *parse_tree_);
+    }
     CARBON_VLOG() << "*** Parse::Tree ***\n" << parse_tree_;
     if (parse_tree_->has_errors()) {
       success_ = false;
@@ -607,8 +629,12 @@ class Driver::CompilationUnit {
     // to wait for code generation.
     consumer_->Flush();
 
-    CARBON_VLOG() << "*** Raw SemIR::File ***\n" << *sem_ir_ << "\n";
+    if (mem_usage_) {
+      mem_usage_->Collect("sem_ir_", *sem_ir_);
+    }
+
     if (options_.dump_raw_sem_ir && IncludeInDumps()) {
+      CARBON_VLOG() << "*** Raw SemIR::File ***\n" << *sem_ir_ << "\n";
       sem_ir_->Print(driver_->output_stream_, options_.builtin_sem_ir);
       if (options_.dump_sem_ir) {
         driver_->output_stream_ << "\n";
@@ -659,11 +685,16 @@ class Driver::CompilationUnit {
 
   // Runs post-compile logic. This is always called, and called after all other
   // actions on the CompilationUnit.
-  auto PostCompile() const -> void {
+  auto PostCompile() -> void {
     if (options_.dump_shared_values && IncludeInDumps()) {
       Yaml::Print(driver_->output_stream_,
                   value_stores_.OutputYaml(input_filename_));
     }
+    if (mem_usage_) {
+      mem_usage_->Collect("value_stores_", value_stores_);
+      Yaml::Print(driver_->output_stream_,
+                  mem_usage_->OutputYaml(input_filename_));
+    }
 
     // The diagnostics consumer must be flushed before compilation artifacts are
     // destructed, because diagnostics can refer to their state.
@@ -773,6 +804,9 @@ class Driver::CompilationUnit {
 
   bool success_ = true;
 
+  // Tracks memory usage of the compile.
+  std::optional<MemUsage> mem_usage_;
+
   // These are initialized as steps are run.
   std::optional<SourceBuffer> source_;
   std::optional<Lex::TokenizedBuffer> tokens_;

toolchain/driver/testdata/dump_mem_usage.carbon

@@ -0,0 +1,19 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// ARGS: compile --phase=check --dump-mem-usage %s
+//
+// NOAUTOUPDATE
+//
+// Avoid testing specific values:
+// SET-CHECK-SUBSET
+
+var x: i32 = 1;
+
+// CHECK:STDOUT: ---
+// CHECK:STDOUT: filename:        dump_mem_usage.carbon
+// CHECK:STDOUT: source_:
+// CHECK:STDOUT:   used_bytes:      0
+// CHECK:STDOUT:   reserved_bytes:  0
+// CHECK:STDOUT: ...

toolchain/lex/BUILD

@@ -221,6 +221,7 @@ cc_library(
         "//common:ostream",
         "//common:string_helpers",
         "//toolchain/base:index_base",
+        "//toolchain/base:mem_usage",
         "//toolchain/base:value_store",
         "//toolchain/diagnostics:diagnostic_emitter",
         "//toolchain/source:source_buffer",

toolchain/lex/tokenized_buffer.cpp

@@ -340,6 +340,13 @@ auto TokenizedBuffer::AddToken(TokenInfo info) -> TokenIndex {
   return TokenIndex(static_cast<int>(token_infos_.size()) - 1);
 }
 
+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_);
+}
+
 auto TokenIterator::Print(llvm::raw_ostream& output) const -> void {
   output << token_.index;
 }

toolchain/lex/tokenized_buffer.h

@@ -18,6 +18,7 @@
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/raw_ostream.h"
 #include "toolchain/base/index_base.h"
+#include "toolchain/base/mem_usage.h"
 #include "toolchain/base/value_store.h"
 #include "toolchain/diagnostics/diagnostic_emitter.h"
 #include "toolchain/lex/token_index.h"
@@ -204,6 +205,10 @@ class TokenizedBuffer : public Printable<TokenizedBuffer> {
   auto PrintToken(llvm::raw_ostream& output_stream, TokenIndex token) const
       -> void;
 
+  // Collects memory usage of members.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void;
+
   // Returns true if the buffer has errors that were detected at lexing time.
   auto has_errors() const -> bool { return has_errors_; }
 

toolchain/parse/tree.cpp

@@ -173,6 +173,12 @@ auto Tree::Print(llvm::raw_ostream& output, bool preorder) const -> void {
   output << "  ]\n";
 }
 
+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_);
+}
+
 auto Tree::VerifyExtract(NodeId node_id, NodeKind kind,
                          ErrorBuilder* trace) const -> bool {
   switch (kind) {

toolchain/parse/tree.h

@@ -232,6 +232,10 @@ class Tree : public Printable<Tree> {
   // line-oriented shell tools from `grep` to `awk`.
   auto Print(llvm::raw_ostream& output, bool preorder) const -> void;
 
+  // Collects memory usage of members.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void;
+
   // The following `Extract*` function provide an alternative way of accessing
   // the nodes of a tree. It is intended to be more convenient and type-safe,
   // but slower and can't be used on nodes that are marked as having an error.

toolchain/sem_ir/block_value_store.h

@@ -81,6 +81,14 @@ class BlockValueStore : public Yaml::Printable<BlockValueStore<IdT>> {
     });
   }
 
+  // Collects memory usage of members.
+  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));
+  }
+
   auto size() const -> int { return values_.size(); }
 
  protected:

toolchain/sem_ir/constant.h

@@ -96,6 +96,14 @@ class ConstantValueStore {
     return symbolic_constants_[const_id.symbolic_index()];
   }
 
+  // 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_);
+  }
+
   // Returns the constant values mapping as an ArrayRef whose keys are
   // instruction indexes. Some of the elements in this mapping may be Invalid or
   // NotConstant.
@@ -133,6 +141,13 @@ class ConstantStore {
   // constant is new.
   auto GetOrAdd(Inst inst, bool is_symbolic) -> ConstantId;
 
+  // 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_);
+  }
+
   // Returns a copy of the constant IDs as a vector, in an arbitrary but
   // stable order. This should not be used anywhere performance-sensitive.
   auto array_ref() const -> llvm::ArrayRef<InstId> { return constants_; }

toolchain/sem_ir/file.cpp

@@ -172,6 +172,32 @@ 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, "bind_names_"), bind_names_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "functions_"), functions_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "classes_"), classes_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "interfaces_"), interfaces_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "impls_"), impls_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "generics_"), generics_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "generic_instances_"),
+                    generic_instances_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "import_irs_"), import_irs_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "import_ir_insts_"),
+                    import_ir_insts_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "type_blocks_"), type_blocks_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "insts_"), insts_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "name_scopes_"), name_scopes_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "constant_values_"),
+                    constant_values_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "inst_blocks_"), inst_blocks_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "constants_"), constants_);
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "types_"), types_);
+  mem_usage.Add(MemUsage::ConcatLabel(label, "complete_types_"),
+                complete_types_);
+}
+
 // Map an instruction kind representing a type into an integer describing the
 // precedence of that type's syntax. Higher numbers correspond to higher
 // precedence.

toolchain/sem_ir/file.h

@@ -53,6 +53,10 @@ class File : public Printable<File> {
   }
   auto OutputYaml(bool include_builtins) const -> Yaml::OutputMapping;
 
+  // Collects memory usage of members.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void;
+
   // Returns array bound value from the bound instruction.
   auto GetArrayBoundValue(InstId bound_id) const -> uint64_t {
     return ints()

toolchain/sem_ir/generic.cpp

@@ -44,6 +44,15 @@ auto GenericInstanceStore::GetOrAdd(GenericId generic_id, InstBlockId args_id)
       .key();
 }
 
+auto GenericInstanceStore::CollectMemUsage(MemUsage& mem_usage,
+                                           llvm::StringRef label) const
+    -> void {
+  mem_usage.Collect(MemUsage::ConcatLabel(label, "generic_instances_"),
+                    generic_instances_);
+  mem_usage.Add(MemUsage::ConcatLabel(label, "lookup_table_"), lookup_table_,
+                KeyContext(generic_instances_.array_ref()));
+}
+
 auto GetConstantInInstance(const File& sem_ir,
                            GenericInstanceId /*instance_id*/,
                            ConstantId const_id) -> ConstantId {

toolchain/sem_ir/generic.h

@@ -120,6 +120,10 @@ class GenericInstanceStore : public Yaml::Printable<GenericInstanceStore> {
     return generic_instances_.OutputYaml();
   }
 
+  // Collects memory usage of members.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void;
+
   auto array_ref() const -> llvm::ArrayRef<GenericInstance> {
     return generic_instances_.array_ref();
   }

toolchain/sem_ir/impl.h

@@ -74,6 +74,13 @@ class ImplStore {
     return values_.OutputYaml();
   }
 
+  // Collects memory usage of members.
+  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_);
+  }
+
   auto array_ref() const -> llvm::ArrayRef<Impl> { return values_.array_ref(); }
   auto size() const -> size_t { return values_.size(); }
 

toolchain/sem_ir/inst.h

@@ -423,6 +423,13 @@ class InstStore {
     values_.Reserve(size);
   }
 
+  // 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, "values_"), values_);
+  }
+
   auto array_ref() const -> llvm::ArrayRef<Inst> { return values_.array_ref(); }
   auto size() const -> int { return values_.size(); }
 

toolchain/sem_ir/name_scope.h

@@ -161,6 +161,12 @@ class NameScopeStore {
     return values_.OutputYaml();
   }
 
+  // Collects memory usage of members.
+  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
+      -> void {
+    mem_usage.Collect(label, values_);
+  }
+
  private:
   InstStore* insts_;
   ValueStore<NameScopeId> values_;

toolchain/sem_ir/type.h

@@ -103,6 +103,11 @@ class TypeStore : public ValueStore<TypeId> {
     return int_type && int_type->int_kind.is_signed();
   }
 
+  // This has no direct storage, so there's no interesting memory usage to
+  // collect.
+  auto CollectMemUsage(MemUsage& /*mem_usage*/, llvm::StringRef /*label*/) const
+      -> void {}
+
  private:
   InstStore* insts_;
   ConstantValueStore* constants_;