M
Modular•7mo 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
benny•7mo 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
tkeittOP•7mo 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
benny•7mo ago
can you give example code?
tkeitt
tkeittOP•7mo 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
benny•7mo ago
you could also use a global variable for config and just import it and use it everywhere

Did you find this page helpful?