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•3mo 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•3mo 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•3mo ago

can you give example code?

tkeitt•3mo 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•3mo ago

you could also use a global variable for config and just import it and use it everywhere

Gaming

Programming

Traits with compile-time SIMD size