Device API Review

[ayerofieiev-tt]

Device API Review

Device class has 177 symbols exposed (members and member functions).
After initial cleanup we get to 132, which is still unreasonable.
Here we share the review of existing Device functions.

Summary

• tt_metal/impl/device/device.hpp include is actively used even in higher layers of the stack
• ~30 methods proxy calls to tt::Cluster
• ~25 methods proxy calls to Allocator
• ~7-10 methods are only used by DevicePool to manage device's lifecycle
• ~5-7 methods used to build programs
• ~5 methods are in support of Trace functionality
• Some methods are duplicated with/without SubDeviceIds argument
• Some methods are only used by dispatch
• Only a subset of methods actively used above Metal (e.g. in TT-NN)

Recommendations

• Review the group of methods related to Cluster
  • Some methods depend on the ARCH, some depend on the actual device (e.g. core_grid vs logical_core_grid)
  • Some methods might be moved out to HAL
  • Some methods might be removed in favor of a single get_soc_descriptor() const
• Review the group of methods related to Allocator
  • Allocator can mature into a proper class/interface, moving some of the functionality from free functions directly into it
  • Some methods might be removed, letting users query info directly from the Allocator
• Align on how to best navigate internal (used in metal, don't mixup with C++ private) vs external (used in layers above metal) interface

Details

These methods are just proxying a call to Cluster giving it device id.
For example

int Device::num_dram_channels() const {
   return tt::Cluster::instance().get_soc_desc(id_).get_num_dram_channels();
}

It is convenient, but it also means MeshDevice and MockDevice have to implement.
Where possible - return metal_SocDescriptor or let user ask Cluser directly

Whats arch specific vs device specific?

   int num_dram_channels() const; X
   uint32_t l1_size_per_core() const;X
   uint32_t dram_size_per_channel() const; X
   CoreCoord grid_size() const; X
   CoreCoord logical_grid_size() const; 
   CoreCoord dram_grid_size() const; X
   CoreType core_type_from_virtual_core(const CoreCoord& virtual_coord) const;
   // Given a Virtual coordinate in noc_index space, get the equivalent coordinate in Virtual NOC0 space
   CoreCoord virtual_noc_coordinate(uint8_t noc_index, CoreCoord coord) const;
   // Given a coordinate in Virtual NOC0 Space, get the equivalent coordinate in Virtual noc_index space
   CoreCoord virtual_noc0_coordinate(uint8_t noc_index, CoreCoord coord) const;

   std::vector<CoreCoord> worker_cores_from_logical_cores(const std::vector<CoreCoord> &logical_cores) const;
   std::vector<CoreCoord> ethernet_cores_from_logical_cores(const std::vector<CoreCoord> &logical_cores) const;
   std::vector<CoreCoord> get_optimal_dram_bank_to_logical_worker_assignment();

   CoreCoord virtual_core_from_logical_core(const CoreCoord &logical_coord, const CoreType& core_type) const;
   CoreCoord worker_core_from_logical_core(const CoreCoord &logical_core) const;

   // Ethernet API
   CoreCoord ethernet_core_from_logical_core(const CoreCoord &logical_core) const;
   CoreCoord logical_core_from_ethernet_core(const CoreCoord &ethernet_core) const;
   std::unordered_set<CoreCoord> get_active_ethernet_cores(bool skip_reserved_tunnel_cores=false) const;
   std::unordered_set<CoreCoord> get_inactive_ethernet_cores() const;
   bool is_active_ethernet_core(CoreCoord logical_core, bool skip_reserved_tunnel_cores=false) const;
   std::tuple<chip_id_t, CoreCoord> get_connected_ethernet_core(CoreCoord eth_core) const;
   std::vector<CoreCoord> get_ethernet_sockets(chip_id_t connected_chip_id) const;
   bool is_inactive_ethernet_core(CoreCoord logical_core) const;

   CoreCoord compute_with_storage_grid_size() const;

   CoreRangeSet worker_cores(HalProgrammableCoreType core_type, SubDeviceId sub_device_id) const;
   uint32_t num_worker_cores(HalProgrammableCoreType core_type, SubDeviceId sub_device_id) const;

   bool is_mmio_capable() const;

   uint32_t num_sub_devices() const;

Next group of methods is proxying a call to Allocator.
For example

DeviceAddr Device::get_base_allocator_addr(const HalMemType &mem_type) const {
   const auto& allocator = this->get_initialized_allocator();
   return allocator::get_unreserved_base_address(*allocator, mem_type);
}

Convert Allocator to a class. Move these convinience methods to it.
Let user get allocator from device and query info from there

   const std::unique_ptr<Allocator> &get_initialized_allocator() const;
   const std::unique_ptr<Allocator> &get_initialized_allocator(SubDeviceId sub_device_id) const;

   DeviceAddr get_base_allocator_addr(const HalMemType &mem_type) const;
   DeviceAddr get_base_allocator_addr(const HalMemType &mem_type, SubDeviceId sub_device_id) const;

   uint32_t num_banks(const BufferType &buffer_type) const;
   uint32_t num_banks(const BufferType &buffer_type, SubDeviceId sub_device_id) const;
   uint32_t bank_size(const BufferType &buffer_type) const;
   uint32_t bank_size(const BufferType &buffer_type, SubDeviceId sub_device_id) const;

   uint32_t dram_channel_from_bank_id(uint32_t bank_id) const;
   uint32_t dram_channel_from_bank_id(uint32_t bank_id, SubDeviceId sub_device_id) const;

   CoreCoord logical_core_from_dram_channel(uint32_t dram_channel) const;
   uint32_t dram_channel_from_logical_core(const CoreCoord& logical_core) const;

   int32_t bank_offset(BufferType buffer_type, uint32_t bank_id) const;
   int32_t bank_offset(BufferType buffer_type, uint32_t bank_id, SubDeviceId sub_device_id) const;

   CoreCoord logical_core_from_bank_id(uint32_t bank_id) const;
   CoreCoord logical_core_from_bank_id(uint32_t bank_id, SubDeviceId sub_device_id) const;

   const std::vector<uint32_t> &bank_ids_from_dram_channel(uint32_t dram_channel) const;
   const std::vector<uint32_t> &bank_ids_from_dram_channel(uint32_t dram_channel, SubDeviceId sub_device_id) const;

   const std::vector<uint32_t> &bank_ids_from_logical_core(
       BufferType buffer_type, const CoreCoord &logical_core) const;
   const std::vector<uint32_t> &bank_ids_from_logical_core(
       BufferType buffer_type, const CoreCoord &logical_core, SubDeviceId sub_device_id) const;

   allocator::Statistics get_memory_allocation_statistics(const BufferType &buffer_type) const;
   allocator::Statistics get_memory_allocation_statistics(const BufferType &buffer_type, SubDeviceId sub_device_id) const;

   uint32_t get_allocator_alignment() const;
   uint32_t get_allocator_alignment(SubDeviceId sub_device_id) const;

   std::optional<DeviceAddr> lowest_occupied_compute_l1_address() const;
   std::optional<DeviceAddr> lowest_occupied_compute_l1_address(tt::stl::Span<const SubDeviceId> sub_device_ids) const;

   size_t get_l1_small_size() const;
   size_t get_l1_small_size(SubDeviceId sub_device_id) const;

   const std::unordered_set<Buffer *> &get_allocated_buffers() const;
   const std::unordered_set<Buffer *> &get_allocated_buffers(SubDeviceId sub_device_id) const;

   void deallocate_buffers();
   void deallocate_buffers(SubDeviceId sub_device_id);

   void dump_memory_blocks(const BufferType &buffer_type, std::ofstream &out) const;
   void dump_memory_blocks(const BufferType &buffer_type, std::ofstream &out, SubDeviceId sub_device_id) const;

   // Set of logical ethernet core coordinates
   // core.x represents connectivity to one other chip, i.e. cores with <x> all connect to same chip
   // core.y represents different channels along one <x>
   const std::set<CoreCoord> &ethernet_cores() const { return this->ethernet_cores_; }

   const std::set<CoreCoord> &storage_only_cores() const { return this->storage_only_cores_; }

   uint32_t get_noc_unicast_encoding(uint8_t noc_index, const CoreCoord& core) const;
   uint32_t get_noc_multicast_encoding(uint8_t noc_index, const CoreRange& cores) const;

   // machine epsilon
   float sfpu_eps() const;

   // machine nan
   float sfpu_nan() const;

   // machine inf
   float sfpu_inf() const;

Only used in tests

   const string build_firmware_target_path(uint32_t programmable_core, uint32_t processor_class, int i) const;

Can be extracted out of device? (No longer has to be this way as of couple weeks ago)

   const JitBuildEnv& build_env() const { return this->build_env_; }
   const string build_kernel_target_path(uint32_t programmable_core, uint32_t processor_class, int i, const string& kernel_name) const;
   const JitBuildState& build_firmware_state(uint32_t programmable_core, uint32_t processor_class, int i) const;
   const JitBuildState& build_kernel_state(uint32_t programmable_core, uint32_t processor_class, int i) const;
   const JitBuildStateSubset build_kernel_states(uint32_t programmable_core, uint32_t processor_class) const;

   SystemMemoryManager& sysmem_manager() { return *sysmem_manager_; }
   HWCommandQueue& hw_command_queue(size_t cq_id = 0);
   CommandQueue& command_queue(size_t cq_id = 0);

   // Metal trace device capture mode
   void begin_trace(const uint8_t cq_id, const uint32_t tid);
   void end_trace(const uint8_t cq_id, const uint32_t tid);
   void replay_trace(const uint8_t cq_id, const uint32_t tid, const bool blocking);
   void release_trace(const uint32_t tid);
   std::shared_ptr<TraceBuffer> get_trace(uint32_t tid);

We don’t really need them both

   bool using_slow_dispatch() const;
   bool using_fast_dispatch() const { return using_fast_dispatch_; }

Next group could be private but they are actively used by DevicePool to maintain device lifecycle

   // Checks that the given arch is on the given pci_slot and that it's responding
   // Puts device into reset
   bool initialize(const uint8_t num_hw_cqs, size_t l1_small_size, size_t trace_region_size, tt::stl::Span<const std::uint32_t> l1_bank_remap = {}, bool minimal = false);
   void build_firmware();
   void reset_cores();
   void initialize_and_launch_firmware();
   void init_command_queue_host();
   void init_command_queue_device();
   void initialize_synchronous_sw_cmd_queue();
   void update_dispatch_cores_for_multi_cq_eth_dispatch();

17 arguments, (already fixed)

   void configure_kernel_variant(Program& program, const string& path, const std::vector<uint32_t>& compile_args, CoreCoord kernel_core, CoreCoord Kernel_virtual_core,
                                 CoreType dispatch_core_type, CoreCoord upstream_virtual_core, CoreCoord downstream_virtual_core, CoreCoord downstream_slave_virtual_core, std::map<string, string> defines_in, NOC my_noc_index, NOC upstream_noc_index, NOC downstream_noc_index, bool is_active_eth_core = false, bool send_to_brisc = false, bool force_watcher_no_inline = false);

   // Puts device into reset
   bool close();
   friend bool CloseDevice(Device *device);

Only used in Buffer

   bool can_use_passthrough_scheduling() const;

Template will have to be removed for interface

   template<typename F>
   void push_work(F&& work, bool blocking = false) {
       this->work_executor_.push_work(std::forward<F>(work), blocking);
   }

   void synchronize();
   void enable_async(bool enable);

We are leaking a detail about Worker here. If we are ok about it - we should remove methods and give worker outside?

   WorkExecutorMode get_worker_mode() { return work_executor_.get_worker_mode(); }
   void set_worker_queue_mode(const WorkerQueueMode& mode) { this->work_executor_.set_worker_queue_mode(mode); }
   WorkerQueueMode get_worker_queue_mode() { return this->work_executor_.get_worker_queue_mode(); }

   // Program cache interface. Syncrhonize with worker worker threads before querying or
   // modifying this structure, since worker threads use this for compiling ops
   void enable_program_cache();
   void disable_and_clear_program_cache();
   program_cache::detail::ProgramCache& get_program_cache() { return program_cache_; }

Delete, users can ask ProgramCache directly

   std::size_t num_program_cache_entries();

   HalProgrammableCoreType get_programmable_core_type(CoreCoord virtual_core) const;

Lets review if we can remove the template from here

   template <typename T = DeviceAddr>
   T get_dev_addr(CoreCoord virtual_core, HalL1MemAddrType addr_type) const;

Only used by Program

   std::vector<std::pair<transfer_info_cores, uint32_t>> extract_dst_noc_multicast_info(const std::vector<CoreRange>& ranges, const CoreType core_type);

Have no idea what these methods are

   bool dispatch_s_enabled() const;
   bool distributed_dispatcher() const;
   NOC dispatch_go_signal_noc() const;
   size_t get_device_kernel_defines_hash();

   uint8_t num_noc_mcast_txns(SubDeviceId sub_device_id) const;
   uint8_t num_noc_unicast_txns(SubDeviceId sub_device_id) const;
   uint8_t noc_data_start_index(SubDeviceId sub_device_id, bool mcast_data=true, bool unicast_data=true) const;

   LaunchMessageRingBufferState& get_worker_launch_message_buffer_state(SubDeviceId sub_device_id);

Subdevice concept and apis must be reviewed and revisited

   SubDeviceManagerId get_active_sub_device_manager_id() const;
   SubDeviceManagerId get_default_sub_device_manager_id() const;
   SubDeviceManagerId create_sub_device_manager(tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size);
   void load_sub_device_manager(SubDeviceManagerId sub_device_manager_id);
   void clear_loaded_sub_device_manager();
   void remove_sub_device_manager(SubDeviceManagerId sub_device_manager_id);
   const std::vector<SubDeviceId> &get_sub_device_ids() const;

   // TODO #15944: Temporary api until migration to actual fabric is complete
   std::tuple<SubDeviceManagerId, SubDeviceId> create_sub_device_manager_with_fabric(tt::stl::Span<const SubDevice> sub_devices, DeviceAddr local_l1_size);
   std::optional<SubDeviceId> get_fabric_sub_device_id() const;

Single usage in Command Queue

   uint32_t get_completion_queue_reader_core() const { return completion_queue_reader_core_; }

Single usage in tt_metal

   bool is_worker_queue_empty() const { return work_executor_.worker_queue.empty(); }

Trace maintains the value and device is only used to store the value, bad design

   uint32_t get_trace_buffers_size() const { return trace_buffers_size_; }
   void set_trace_buffers_size(uint32_t size) { trace_buffers_size_ = size; }

Single use in device pool

   std::vector<std::vector<chip_id_t>> get_tunnels_from_mmio() const { return tunnels_from_mmio_; }

Should device init allocator wit this scheme and should users ask allocator directly?

   static constexpr MemoryAllocator allocator_scheme_ = MemoryAllocator::L1_BANKING;

[ayerofieiev-tt]

#16482 (comment)

CC @tt-asaigal