M
Modular5mo ago
tkeitt

Traits with compile-time SIMD size

I can't seem to fulfill a trait that can return different length SIMDs
trait PRNGEngine(Movable):
@staticmethod
fn ndim() -> Int:
"""SIMD size."""
pass

fn next_scalar(inout self) -> UInt64:
"""Get only a single value from the generator."""
pass

fn next(inout self) -> SIMD[DType.uint64, Self.ndim()]:
"""Get Self.ndim() values from the generator."""
pass


struct DummyEngine[n: Int](PRNGEngine):
@staticmethod
fn ndim() -> Int:
"""SIMD size."""
return 2

fn next_scalar(inout self) -> UInt64:
"""Get only a single value from the generator."""
return 42

fn next(inout self) -> SIMD[DType.uint64, Self.ndim()]:
"""Get Self.ndim() values from the generator."""
var res: SIMD[DType.uint64, Self.ndim()] = 42
return res
trait PRNGEngine(Movable):
@staticmethod
fn ndim() -> Int:
"""SIMD size."""
pass

fn next_scalar(inout self) -> UInt64:
"""Get only a single value from the generator."""
pass

fn next(inout self) -> SIMD[DType.uint64, Self.ndim()]:
"""Get Self.ndim() values from the generator."""
pass


struct DummyEngine[n: Int](PRNGEngine):
@staticmethod
fn ndim() -> Int:
"""SIMD size."""
return 2

fn next_scalar(inout self) -> UInt64:
"""Get only a single value from the generator."""
return 42

fn next(inout self) -> SIMD[DType.uint64, Self.ndim()]:
"""Get Self.ndim() values from the generator."""
var res: SIMD[DType.uint64, Self.ndim()] = 42
return res
I can fall back to next_scalar. I'd like to be able to use next e.g. to generate a set of normal deviates, etc.
5 Replies
benny
benny5mo ago
trait PRNGEngine(Movable):
fn next_scalar(inout self) -> UInt64:
...

fn next[W: Int](inout self) -> SIMD[DType.uint64, W]:
...

struct DummyEngine(PRNGEngine):
fn next_scalar(inout self) -> UInt64:
return 42

fn next[W: Int](inout self) -> SIMD[DType.uint64, W]:
return SIMD[DType.uint64, W](42)
trait PRNGEngine(Movable):
fn next_scalar(inout self) -> UInt64:
...

fn next[W: Int](inout self) -> SIMD[DType.uint64, W]:
...

struct DummyEngine(PRNGEngine):
fn next_scalar(inout self) -> UInt64:
return 42

fn next[W: Int](inout self) -> SIMD[DType.uint64, W]:
return SIMD[DType.uint64, W](42)
does that fix the issue? If you want it to be static for all the calls and declared for the struct itself you could altnernatively do
trait PRNGEngine(Movable):
fn next_scalar(inout self) -> UInt64:
...

fn next(inout self) -> SIMD[DType.uint64, W]:
...

struct DummyEngine[W: Int](PRNGEngine):
fn next_scalar(inout self) -> UInt64:
return 42

fn next(inout self) -> SIMD[DType.uint64, W]:
return SIMD[DType.uint64, W](42)
trait PRNGEngine(Movable):
fn next_scalar(inout self) -> UInt64:
...

fn next(inout self) -> SIMD[DType.uint64, W]:
...

struct DummyEngine[W: Int](PRNGEngine):
fn next_scalar(inout self) -> UInt64:
return 42

fn next(inout self) -> SIMD[DType.uint64, W]:
return SIMD[DType.uint64, W](42)
NOTE: There are pros and cons to both of these With the first approach, W must be specified every time .next is called, or given a default value With the second approach, DummyEngine can't be easily stored in a collection, since DummyEngine[1] and DummyEngine[2] are different types
tkeitt
tkeittOP5mo ago
I tried that a bunch of different ways. When I parameterize PRNG as PRNG[T: PRNGEngine], I can't seem to resolve W.
benny
benny5mo ago
can you give example code?
tkeitt
tkeittOP5mo ago
I'll work something up that is a bit more complete. I am trying to infer W when passing the engine into the PRNG. I can see that it will probably work if I write it the dimensions explicitly everywhere.
benny
benny5mo ago
you could also use a global variable for config and just import it and use it everywhere
Want results from more Discord servers?
Add your server