Modular•10mo ago

How to rewrite this code into something not ugly

I need to call shuffle on a parameter mask of different lengths, the following code is the shortest that I could make. Please fill in the dots to appreciate what would happen with width 1024. Any thoughts are appreciated.

fn my_shuffle[T: DType, width: Int, p: StaticIntTuple[width]](v: SIMD[T, width]) -> SIMD[T, width]:
    @parameter
    if width == 8:
        return v.shuffle[
            p[0],
            p[1],
            p[2],
            p[3],
            p[4],
            p[5],
            p[6],
            p[7],
        ]()
    elif width == 16:
        return v.shuffle[
            p[0],
...
            p[15],
        ]()
    elif width == 32:
        return v.shuffle[
            p[0],
...
            p[31],
        ]()
    elif width == 64:
        return v.shuffle[
            p[0],
...   
            p[63],
        ]()
    elif width == 128:
        return v.shuffle[
            p[0],
...
            p[127],
        ]()
    elif width == 256:
        return v.shuffle[
            p[0],
...
            p[255],
        ]()
    else:
        constrained[False]()
        return v

fn my_shuffle[T: DType, width: Int, p: StaticIntTuple[width]](v: SIMD[T, width]) -> SIMD[T, width]:
    @parameter
    if width == 8:
        return v.shuffle[
            p[0],
            p[1],
            p[2],
            p[3],
            p[4],
            p[5],
            p[6],
            p[7],
        ]()
    elif width == 16:
        return v.shuffle[
            p[0],
...
            p[15],
        ]()
    elif width == 32:
        return v.shuffle[
            p[0],
...
            p[31],
        ]()
    elif width == 64:
        return v.shuffle[
            p[0],
...   
            p[63],
        ]()
    elif width == 128:
        return v.shuffle[
            p[0],
...
            p[127],
        ]()
    elif width == 256:
        return v.shuffle[
            p[0],
...
            p[255],
        ]()
    else:
        constrained[False]()
        return v

9 Replies

Michael K•10mo ago

I have written this code as well, using shuffle to do deinterleave since that functions is broken. I think some day there will be sugar to unpack something you can slice. Or for shuffleto take a parameter that can be sliced. But probably not a priority.

sora•10mo ago

Maybe you could use v._shuffle_list as a workaround for now?

Henk-Jan LebbinkOP•10mo ago

v._shuffle_list is an alternative, but it also has its problems. E.g., I don't know how to fill a VariadicList other than with the following method, which does not solve the inconvenience of enumerating all elements of the mask.

  @parameter
  if width == 8:
      alias x = VariadicList[Int](p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7])
      return v._shuffle_list[x, width](v)
  elif width == 16:
      alias x = ....

  @parameter
  if width == 8:
      alias x = VariadicList[Int](p[0], p[1], p[2], p[3], p[4], p[5], p[6], p[7])
      return v._shuffle_list[x, width](v)
  elif width == 16:
      alias x = ....

sora•10mo ago

Yea, I also tried it and it's horrible.

Henk-Jan LebbinkOP•10mo ago

Oh well, it compiles at least into something good.

sora•10mo ago

Maybe we should ask how to turn a StaticTuple into a VariadicList. There must be a way, as solution to such a problem is recently add to polynomial_evaluate. See change log of 0.7, also GitHub issue 1587.

Henk-Jan LebbinkOP•10mo ago

Someone said that VariadicList was not intended to be created and passed around (but I can't find that quote anymore). Yes, lets ask for that, or even better to update the signatures of the shuffle methods.

sora•10mo ago

Yea, I also remember that. VariadicList does look like a language<->data model interface type. I would expect to write the code someday simply as v.shuffle[*v]().

Michael K•8mo ago

I opened an issue for this.

GitHub

[Feature Request] Allow StaticIntTuple mask for SIMD.shuffle · Issu...

Review Mojo's priorities I have read the roadmap and priorities and I believe this request falls within the priorities. What is your request? As title. I want to build the list of indices in th...

Gaming

Programming

How to rewrite this code into something not ugly