Remix.run Logo
LegNeato 22 minutes ago

Didn't want to go into crazy detail in the post.

Each family of operations is a trait parameterized by the operation itself:

  pub trait EvaluateReduction<Operation, T>: LaneEvaluator {
      /// Reduce one distributed definition to an ordinary uniform scalar.
      fn evaluate_reduction(&self, value: LaneValue<Self, role::Distributed, T>) -> T;
  }

Call sites name the operation:

  let one   = evaluator.splat::<Splat, _>(1_u32);
  let two   = evaluator.splat::<Splat, _>(2_u32);
  let three = evaluator.binary::<Add, _>(one, two);

  let total   = evaluator.reduce::<Sum, u32>(three);   // a uniform u32
  let running = <Executor as EvaluateScan<Scan<Sum, Exclusive>, u32>>::scan(&evaluator, three);

Operations like Sum, Max, ReduceXor, Inclusive, and Exclusive are all distinct types.

As mentioned in the post, execution shape is typed too. A static shuffle takes its control as a type-level constant, and the shuffle mode constrains which controls are expressible:

  // Shift down one lane, keeping our own value where the source is inactive.
  let down  = <Executor as EvaluateShuffle<Shuffle<Down>, DownOrSelf<1>, u32>>::shuffle(&ev, v);
  // Broadcast from lane zero.
  let bcast = <Executor as EvaluateShuffle<Shuffle<Broadcast>, WarpLane<0>, u32>>::shuffle(&ev, down);
  // Butterfly exchange with the neighbor one bit away.
  let bfly  = <Executor as EvaluateShuffle<Shuffle<Xor>, Butterfly<1>, u32>>::shuffle(&ev, bcast);

For an example of errors caught, a warp-scoped executor for a device-scoped barrier is a compile error:

  <ScopedWarpExecutor<'_, WarpUniform> as EvaluateBarrier<Barrier<Device>>>::barrier(evaluator)
  // error[E0277]: the trait bound `Device: NvptxBarrierScope` is not satisfied
  //               help: the trait `NvptxBarrierScope` is implemented for `Warp`

Strip mining is typed on the amount of work and the lane capacity, and it hands back one chunk at a time along with the predicate saying which lanes live in that chunk:

  // Six work items across four active lanes: two chunks, based at 0 and 4.
  <Executor as EvaluateStripMine<StripMine, (WorkItems, ActiveLanes<StripMined<4>>), i32>>::
      for_each_strip_mined(
          &evaluator,
          (WorkItems::new(6)?, ActiveLanes::new(4)?),
          |index, active| {
           // ...
          },
      );

Hopefully that gives the flavor of it.