One thing I have repeatedly run into is that Mojo’s capabilities for abstraction at runtime feel inferior to Mojo’s compile-time abstraction capabilties. While this is often fine for hot loops and fixable with the JIT compiler, there are numerous other places where this would require JITing the entire program.
Common Features That Desire Runtime Abstraction
Logging backends
stderr, stderr with colors, stderr as jsonl, syslog, a single file, rotating files in a directory, open telemetry, sentry, and more, there are a lot of ways that people want to log things and you often want the ability to pick how logging is done at runtime.
GPU Telemetry
It can be useful to report if a GPU is spending substantial amounts of time with its PCIe interface maxed out, or if some of its vram is already in use by another process, such as a local user’s desktop and browser. A more industrial system such as Mammoth might want to monitor NVLink error rates on NVIDIA, while handling UALink on AMD with the same binary. While MAX could be used to JIT compile the entire telemetry module of a codebase, this seems like overkill.
Optional GPU offloading
Due to the lack of a DeviceContext API to access CPU acceleration, there will be a level of runtime abstraction if a program wishes to cleanly fall back to CPU if no GPU is present. This is a very important feature for software which cannot guarantee users have a supported GPU, and as a result much of the software in this category is currently stuck with OpenCL, making Mojo a massive upgrade.
Abstractions over MAX and other APIs
OpenCL, Vulkan and DirectX Compute (once Mojo supports Windows) can be used to provide a fallback if MAX does not support the devices present on the user’s system.
Allocators chosen at runtime
While an allocator that isn’t NUMA aware is a performance benefit on a monolithic consumer chip, one which handles NUCA (non-uniform cache architecture) is quite helpful on multi-CCD AMD, and of course we want the binary to handle NUMA when it’s on a full NUMA system. Additionally, since tcmalloc requires knowing the page size at compile time, having multiple versions of that around on ARM would be helpful.
Current State in Mojo
I attempted to write a proper version, as opposed to many of my current once where I threw memory management to C code or just leaked the memory. In the processor of doing that, I ran into this issue. Working around that, I get to the following for a simple application that uses some mocked telemetry providers
from std.logger import Logger
from std.sys.intrinsics import size_of
from std.time import perf_counter_ns
trait TelemetryObject(Copyable, ImplicitlyDestructible, Movable):
comptime WriterType: Writer
def write_as_json_object(self, mut writer: Self.WriterType) raises:
...
struct DynTelemetryObject[WriterTypeInject: Writer](
Copyable, Movable, TelemetryObject
):
comptime WriterType = Self.WriterTypeInject
var object_ptr: OpaquePointer[MutExternalOrigin]
var write_as_json_object_fn: def(
p: OpaquePointer[MutExternalOrigin], mut writer: Self.WriterType
) thin raises
var dtor: def(var p: OpaquePointer[MutExternalOrigin]) thin
def __init__[BaseT: TelemetryObject](out self, var base: BaseT):
def write_as_json_object_fn(
p: OpaquePointer[MutExternalOrigin], mut writer: Self.WriterType
) raises:
var ptr = p.bitcast[BaseT]()
ptr[].write_as_json_object(rebind[BaseT.WriterType](writer))
def dtor(var p: OpaquePointer[MutExternalOrigin]):
var ptr = p.bitcast[BaseT]()
ptr.destroy_pointee()
self.object_ptr = alloc[type_of(base)](1).bitcast[NoneType]()
self.object_ptr.bitcast[BaseT]()[] = base^
self.write_as_json_object_fn = write_as_json_object_fn
self.dtor = dtor
def write_as_json_object(self, mut writer: Self.WriterType) raises:
self.write_as_json_object_fn(self.object_ptr, writer)
def __del__(deinit self):
self.dtor(self.object_ptr)
trait TelemetryProvider(ImplicitlyDestructible, Movable):
comptime TelemetryObjectType: TelemetryObject
def get_telemetry_data(self) -> Self.TelemetryObjectType:
...
struct DynTelemetryProvider[WriterTypeInject: Writer](
Copyable, Movable, TelemetryProvider
):
comptime TelemetryObjectType = DynTelemetryObject[Self.WriterTypeInject]
var provider_ptr: OpaquePointer[MutExternalOrigin]
var get_telemetry_data_fn: def(
p: OpaquePointer[MutExternalOrigin]
) thin -> Self.TelemetryObjectType
var dtor: def(var p: OpaquePointer[MutExternalOrigin]) thin
def __init__[BaseT: TelemetryProvider](out self, var base: BaseT):
self.provider_ptr = alloc[UInt8](size_of[BaseT]()).bitcast[NoneType]()
self.provider_ptr.bitcast[BaseT]()[] = base^
def get_telemetry_data_fn(
p: OpaquePointer[MutExternalOrigin],
) -> Self.TelemetryObjectType:
var ptr = p.bitcast[BaseT]()
return DynTelemetryObject[Self.WriterTypeInject](
ptr[].get_telemetry_data()
)
def dtor(var p: OpaquePointer[MutExternalOrigin]):
var ptr = p.bitcast[BaseT]()
ptr.destroy_pointee()
self.get_telemetry_data_fn = get_telemetry_data_fn
self.dtor = dtor
def get_telemetry_data(self) -> Self.TelemetryObjectType:
return self.get_telemetry_data_fn(self.provider_ptr)
def __del__(deinit self):
self.dtor(self.provider_ptr)
@fieldwise_init
struct NVLinkSwitchTelemetryObject[WriterTypeInject: Writer](
Copyable, Movable, TelemetryObject
):
comptime WriterType = Self.WriterTypeInject
var switch_pcie_id: UInt32
var crc_error_count: UInt64
var link_integrity_errors: UInt64
def write_as_json_object(self, mut writer: Self.WriterType) raises:
writer.write_string(
"{"
+ '"event_name": "nv_link_switch", '
+ ('"switch_pcie_id": ' + String(self.switch_pcie_id) + ", ")
+ ('"crc_error_count": ' + String(self.crc_error_count) + ", ")
+ ('"link_integrity_errors": ' + String(self.link_integrity_errors))
+ "}"
)
@fieldwise_init
struct NVLinkSwitchTelemetryManager[WriterTypeInject: Writer](
Copyable, ImplicitlyDestructible, Movable, TelemetryProvider
):
comptime TelemetryObjectType = NVLinkSwitchTelemetryObject[
Self.WriterTypeInject
]
var switch_pcie_id: UInt32
def get_crc_error_count(self) -> UInt64:
"""Get the CRC error count since startup for this switch."""
return 0 # Placeholder implementation
def get_link_integrity_errors(self) -> UInt64:
"""Get the link integrity error count since startup for this switch."""
return 0 # Placeholder implementation
def get_telemetry_data(
self,
) -> NVLinkSwitchTelemetryObject[Self.WriterTypeInject]:
"""Get the telemetry data for this switch."""
return NVLinkSwitchTelemetryObject[Self.WriterTypeInject](
switch_pcie_id=self.switch_pcie_id,
crc_error_count=self.get_crc_error_count(),
link_integrity_errors=self.get_link_integrity_errors(),
)
@fieldwise_init
struct UALinkSwitchTelemetryObject[WriterTypeInject: Writer](
Copyable, Movable, TelemetryObject
):
comptime WriterType = Self.WriterTypeInject
var switch_pcie_id: UInt32
var crc_error_count: UInt64
var flit_replay_count: UInt64
var link_level_retry_count: UInt64
def write_as_json_object(self, mut writer: Self.WriterType) raises:
writer.write_string(
"{"
+ '"event_name": "ua_link_switch", '
+ ('"switch_pcie_id": ' + String(self.switch_pcie_id) + ", ")
+ ('"crc_error_count": ' + String(self.crc_error_count) + ", ")
+ ('"flit_replay_count": ' + String(self.flit_replay_count) + ", ")
+ (
'"link_level_retry_count": '
+ String(self.link_level_retry_count)
)
+ "}"
)
@fieldwise_init
struct UALinkSwitchTelemetryManager[WriterTypeInject: Writer](
Copyable, Movable, TelemetryProvider
):
comptime TelemetryObjectType = UALinkSwitchTelemetryObject[
Self.WriterTypeInject
]
var switch_pcie_id: UInt32
def get_crc_error_count(self) -> UInt64:
"""Get the CRC error count since startup for this switch."""
return 0 # Placeholder implementation
def get_flit_replay_count(self) -> UInt64:
"""Get the flit replay count since startup for this switch."""
return 0 # Placeholder implementation
def get_link_level_retry_count(self) -> UInt64:
"""Get the link level retry count since startup for this switch."""
return 0 # Placeholder implementation
def get_telemetry_data(
self,
) -> UALinkSwitchTelemetryObject[Self.WriterTypeInject]:
"""Get the telemetry data for this switch."""
return UALinkSwitchTelemetryObject[Self.WriterTypeInject](
switch_pcie_id=self.switch_pcie_id,
crc_error_count=self.get_crc_error_count(),
flit_replay_count=self.get_flit_replay_count(),
link_level_retry_count=self.get_link_level_retry_count(),
)
@fieldwise_init
struct HeartbeatTelemetryObject[WriterTypeInject: Writer](
Copyable, Movable, TelemetryObject
):
comptime WriterType = Self.WriterTypeInject
var timestamp: UInt
def write_as_json_object(self, mut writer: Self.WriterType) raises:
writer.write_string(
"{"
+ '"event_name": "heartbeat", '
+ ('"timestamp": ' + String(self.timestamp))
+ "}"
)
@fieldwise_init
struct HeartbeatTelemetryManager[WriterTypeInject: Writer](
Copyable, Movable, TelemetryProvider
):
comptime TelemetryObjectType = HeartbeatTelemetryObject[
Self.WriterTypeInject
]
def get_timestamp(self) -> UInt:
"""Get the current timestamp in nanoseconds."""
return perf_counter_ns()
def get_telemetry_data(
self,
) -> HeartbeatTelemetryObject[Self.WriterTypeInject]:
"""Get the telemetry data for this heartbeat."""
return HeartbeatTelemetryObject[Self.WriterTypeInject](
self.get_timestamp()
)
@fieldwise_init
struct JsonLogWriter(Movable, Writer):
var logger: Logger[]
def __init__(out self):
self.logger = Logger[]()
def write_string(self, s: StringSlice):
self.logger.info(s)
def write[*Ts: Writable](mut self, *args: *Ts):
var s = String("")
comptime for i in range(len(Ts)):
args[i].write_to(s)
self.write_string(s)
def main() raises:
var logger = JsonLogWriter()
var telemetry_sources: List[DynTelemetryProvider[JsonLogWriter]] = [
DynTelemetryProvider[JsonLogWriter](
NVLinkSwitchTelemetryManager[JsonLogWriter](switch_pcie_id=1)
),
DynTelemetryProvider[JsonLogWriter](
UALinkSwitchTelemetryManager[JsonLogWriter](switch_pcie_id=2)
),
DynTelemetryProvider[JsonLogWriter](
HeartbeatTelemetryManager[JsonLogWriter]()
),
]
for i in range(len(telemetry_sources)):
var telemetry_data = telemetry_sources[i].get_telemetry_data()
telemetry_data.write_as_json_object(logger)
And here’s the much nicer Rust equivalent:
use std::fmt;
// ---------------------------------------------------------------------------
// JsonLogWriter – a simple writer that prints JSON fragments to stdout
// ---------------------------------------------------------------------------
struct JsonLogWriter;
impl JsonLogWriter {
fn write_string(&self, s: &str) {
print!("{}", s);
}
}
// ---------------------------------------------------------------------------
// TelemetryObject – the "data" trait (what gets serialised)
// ---------------------------------------------------------------------------
trait TelemetryObject: fmt::Debug {
fn write_as_json_object(&self, writer: &mut JsonLogWriter);
}
// ---------------------------------------------------------------------------
// TelemetryProvider – the "manager" trait (produces TelemetryObject)
// ---------------------------------------------------------------------------
trait TelemetryProvider: fmt::Debug {
fn get_telemetry_data(&self) -> Box<dyn TelemetryObject>;
}
// ---------------------------------------------------------------------------
// NVLinkSwitch telemetry
// ---------------------------------------------------------------------------
#[derive(Debug)]
struct NVLinkSwitchTelemetryObject {
switch_pcie_id: u32,
crc_error_count: u64,
link_integrity_errors: u64,
}
impl TelemetryObject for NVLinkSwitchTelemetryObject {
fn write_as_json_object(&self, writer: &mut JsonLogWriter) {
writer.write_string(&format!(
"{{\
\"event_name\": \"nv_link_switch\", \
\"switch_pcie_id\": {}, \
\"crc_error_count\": {}, \
\"link_integrity_errors\": {}\
}}",
self.switch_pcie_id, self.crc_error_count, self.link_integrity_errors,
));
}
}
#[derive(Debug)]
struct NVLinkSwitchTelemetryManager {
switch_pcie_id: u32,
}
impl TelemetryProvider for NVLinkSwitchTelemetryManager {
fn get_telemetry_data(&self) -> Box<dyn TelemetryObject> {
Box::new(NVLinkSwitchTelemetryObject {
switch_pcie_id: self.switch_pcie_id,
crc_error_count: 0, // placeholder
link_integrity_errors: 0, // placeholder
})
}
}
// ---------------------------------------------------------------------------
// UALink telemetry
// ---------------------------------------------------------------------------
#[derive(Debug)]
struct UALinkSwitchTelemetryObject {
switch_pcie_id: u32,
crc_error_count: u64,
flit_replay_count: u64,
link_level_retry_count: u64,
}
impl TelemetryObject for UALinkSwitchTelemetryObject {
fn write_as_json_object(&self, writer: &mut JsonLogWriter) {
writer.write_string(&format!(
"{{\
\"event_name\": \"ua_link_switch\", \
\"switch_pcie_id\": {}, \
\"crc_error_count\": {}, \
\"flit_replay_count\": {}, \
\"link_level_retry_count\": {}\
}}",
self.switch_pcie_id,
self.crc_error_count,
self.flit_replay_count,
self.link_level_retry_count,
));
}
}
#[derive(Debug)]
struct UALinkSwitchTelemetryManager {
switch_pcie_id: u32,
}
impl TelemetryProvider for UALinkSwitchTelemetryManager {
fn get_telemetry_data(&self) -> Box<dyn TelemetryObject> {
Box::new(UALinkSwitchTelemetryObject {
switch_pcie_id: self.switch_pcie_id,
crc_error_count: 0, // placeholder
flit_replay_count: 0, // placeholder
link_level_retry_count: 0, // placeholder
})
}
}
// ---------------------------------------------------------------------------
// Heartbeat telemetry
// ---------------------------------------------------------------------------
#[derive(Debug)]
struct HeartbeatTelemetryObject {
timestamp: u64,
}
impl TelemetryObject for HeartbeatTelemetryObject {
fn write_as_json_object(&self, writer: &mut JsonLogWriter) {
writer.write_string(&format!(
"{{\
\"event_name\": \"heartbeat\", \
\"timestamp\": {}\
}}",
self.timestamp,
));
}
}
#[derive(Debug)]
struct HeartbeatTelemetryManager;
impl TelemetryProvider for HeartbeatTelemetryManager {
fn get_telemetry_data(&self) -> Box<dyn TelemetryObject> {
Box::new(HeartbeatTelemetryObject {
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos() as u64,
})
}
}
// ---------------------------------------------------------------------------
// main – demonstrates heterogeneous collection via `Box<dyn TelemetryProvider>`
// ---------------------------------------------------------------------------
fn main() {
let mut writer = JsonLogWriter;
let telemetry_sources: Vec<Box<dyn TelemetryProvider>> = vec![
Box::new(NVLinkSwitchTelemetryManager { switch_pcie_id: 1 }),
Box::new(UALinkSwitchTelemetryManager { switch_pcie_id: 2 }),
Box::new(HeartbeatTelemetryManager),
];
for provider in &telemetry_sources {
let telemetry_data = provider.get_telemetry_data();
telemetry_data.write_as_json_object(&mut writer);
println!();
}
}
This is a substantial amount of advanced Mojo features, along with running into a compiler bug, all to implement what may people would consider a fairly simple bit of functionality. Rust’s version, using trait objects (also known as Existentials), or a C++ version using a bit of inheritance are much easier to implement, so much so that the Rust version was written by a 1.2B parameter LLM from my Mojo code.
Potential Solutions
I don’t think we need full Existentials, just a bare minimum version of “please give me a vtable for that trait who’s functions don’t have parameters” which can be combined with some boilerplate similar to what I have here, or some creative uses of __getattr__, to enable something which behaves very close to Rust’s trait objects.
Reflection-based
In order to solve this with reflection, I would need:
- The ability to create a struct (or modify an existing one)
- The ability to make add a function to a struct
- The ability to make a struct implement a trait
- A way to get the function reflection objects for each function in a trait.
This would be my preferred solution, although it wouldn’t work for this example initially due to not being able to support type aliases, since it provides capabilities that would be extremely useful elsewhere, such as in a gRPC implementation for Mojo.
Compiler Magic
comptime FooVtable = __make_vtable_for(Foo) is fine, although it doesn’t provide additional capabilities for use elsewhere.
Early Classes
Having classes with inheritance also resolves this problem, although those are currently planned for much later.