How to use DeviceContext

GPU Programming
1

If DeviceContext is used in multiple places, is it better to wrap it in ArcPointer and use it via the pointer?

For example:

    var device_ctx1 = DeviceContext()
    var device_ctx2 = DeviceContext()

    # Should we do the following instead? 
    var device_ctx = ArcPointer(DeviceContext())
```

Thanks.
2

The DeviceContext type is already internally ref counted, so you can pass it around in Mojo by value.

Looking at the struct definition in device_context.mojo:

struct DeviceContext(ImplicitlyCopyable, RegisterPassable):
    var _handle: _DeviceContextPtr[mut=True]
    var _owning: Bool

A DeviceContext is a thin Mojo wrapper around a C++ handle that manages its own reference count. The copy constructor calls AsyncRT_DeviceContext_retain and the destructor calls AsyncRT_DeviceContext_release. That means each copy shares the same underlying GPU stream/context, and resources are freed when the last copy goes out of scope - exactly what ArcPointer would give you, but built in.

There are some options based on your example though:

Option 1 (two separate DeviceContext() calls) is most likely not what you want. Each call creates a distinct context with its own stream, so work enqueued on one is independent of the other. That’s only appropriate if you deliberately want concurrent streams.

Option 2 (ArcPointer(DeviceContext())) works but is redundant — you’d be refcounting a refcounted handle.

Idiomatic usage is to create one instance and pass it by value:

def helper(ctx: DeviceContext) raises:
    var buf = ctx.enqueue_create_buffer[DType.float32](N)
…

def main() raises:
    var ctx = DeviceContext()
    helper(ctx)   # implicit copy, cheap refcount bump
    helper(ctx)

You can see this pattern throughout the examples and tests — e.g. mojo/examples/gpu-functions/vector_addition.mojo and mojo/stdlib/test/gpu/host/test_device_context.mojo.

By default DeviceContext() selects device 0 on the default accelerator API for your build target. If you need a second, physically distinct device, pass DeviceContext(device_id=1).

3

Thanks for the explanation - much appreciate it! “helper(ctx)” - pattern works in limited context. If I wanted to throw around an instantiated DeviceContext - embedding the context inside ArcPointer extends the ‘helper’ pattern but incurs the cost of redundant ref counting. My requirement - as you must have understood - how do I throw around a default instantiated DeviceContext in absence of statics/global! "global_constant()" - would not work - would it?

4

You can throw it around by value. Treat it as if it were an object of type Int. The internal ref counting will handle things correctly.

5

I got that. I would have to carry it along all throughout - since we don’t have static/globals. “global_constant()” would not help - am I correct?

6

Correct - you will need to pass it into any function that will need it.