Faster quaternion multiplication

[brentyi]

For pose graph optimization on the sphere2500 dataset, this speeds up overall JIT time by ~35%, from >7 seconds on my machine to 4.5 seconds. Runtime difference is insignificant.

When multiplying quaternions with batch axis (4096,), here's the jaxpr now:

{ lambda a:f32[4096,4] b:f32[4096,4]; . let
    c:f32[4096,4] = pjit[
      name=multiply
      jaxpr={ lambda d:i32[4,4] e:i32[4,4] f:i32[4,4]; g:f32[4096,4] h:f32[4096,4]. let
          i:f32[4096,4,4] = dot_general[
            dimension_numbers=(([], []), ([0], [0]))
            preferred_element_type=float32
          ] g h
          j:bool[4,4] = lt d 0
          k:i32[4,4] = add d 4
          l:i32[4,4] = select_n j d k
          m:bool[4,4] = lt e 0
          n:i32[4,4] = add e 4
          o:i32[4,4] = select_n m e n
          p:i32[4,4,1] = broadcast_in_dim[
            broadcast_dimensions=(0, 1)
            shape=(4, 4, 1)
          ] l
          q:i32[4,4,1] = broadcast_in_dim[
            broadcast_dimensions=(0, 1)
            shape=(4, 4, 1)
          ] o
          r:i32[4,4,2] = concatenate[dimension=2] p q
          s:f32[4096,4,4] = gather[
            dimension_numbers=GatherDimensionNumbers(offset_dims=(0,), collapsed_slice_dims=(1, 2), start_index_map=(1, 2), operand_batching_dims=(), start_indices_batching_dims=())
            fill_value=None
            indices_are_sorted=False
            mode=GatherScatterMode.PROMISE_IN_BOUNDS
            slice_sizes=(4096, 1, 1)
            unique_indices=False
          ] i r
          t:f32[4,4] = convert_element_type[new_dtype=float32 weak_type=False] f
          u:f32[1,4,4] = broadcast_in_dim[
            broadcast_dimensions=(1, 2)
            shape=(1, 4, 4)
          ] t
          v:f32[4096,4,4] = mul u s
          w:f32[4096,4] = reduce_sum[axes=(2,)] v
        in (w,) }
    ] a b
  in (c,) }

Compared to before:

{ lambda a:f32[4096,4] b:f32[4096,4]; . let
    c:f32[4096,4] = pjit[
      name=multiply
      jaxpr={ lambda ; d:f32[4096,4] e:f32[4096,4]. let
          f:f32[4,4096] = transpose[permutation=(1, 0)] d
          g:f32[1,4096] = slice[
            limit_indices=(1, 4096)
            start_indices=(0, 0)
            strides=(1, 1)
          ] f
          h:f32[4096] = squeeze[dimensions=(0,)] g
          i:f32[1,4096] = slice[
            limit_indices=(2, 4096)
            start_indices=(1, 0)
            strides=(1, 1)
          ] f
          j:f32[4096] = squeeze[dimensions=(0,)] i
          k:f32[1,4096] = slice[
            limit_indices=(3, 4096)
            start_indices=(2, 0)
            strides=(1, 1)
          ] f
          l:f32[4096] = squeeze[dimensions=(0,)] k
          m:f32[1,4096] = slice[
            limit_indices=(4, 4096)
            start_indices=(3, 0)
            strides=(1, 1)
          ] f
          n:f32[4096] = squeeze[dimensions=(0,)] m
          o:f32[4,4096] = transpose[permutation=(1, 0)] e
          p:f32[1,4096] = slice[
            limit_indices=(1, 4096)
            start_indices=(0, 0)
            strides=(1, 1)
          ] o
          q:f32[4096] = squeeze[dimensions=(0,)] p
          r:f32[1,4096] = slice[
            limit_indices=(2, 4096)
            start_indices=(1, 0)
            strides=(1, 1)
          ] o
          s:f32[4096] = squeeze[dimensions=(0,)] r
          t:f32[1,4096] = slice[
            limit_indices=(3, 4096)
            start_indices=(2, 0)
            strides=(1, 1)
          ] o
          u:f32[4096] = squeeze[dimensions=(0,)] t
          v:f32[1,4096] = slice[
            limit_indices=(4, 4096)
            start_indices=(3, 0)
            strides=(1, 1)
          ] o
          w:f32[4096] = squeeze[dimensions=(0,)] v
          x:f32[4096] = neg j
          y:f32[4096] = mul x s
          z:f32[4096] = mul l u
          ba:f32[4096] = sub y z
          bb:f32[4096] = mul n w
          bc:f32[4096] = sub ba bb
          bd:f32[4096] = mul h q
          be:f32[4096] = add bc bd
          bf:f32[4096] = mul j q
          bg:f32[4096] = mul l w
          bh:f32[4096] = add bf bg
          bi:f32[4096] = mul n u
          bj:f32[4096] = sub bh bi
          bk:f32[4096] = mul h s
          bl:f32[4096] = add bj bk
          bm:f32[4096] = neg j
          bn:f32[4096] = mul bm w
          bo:f32[4096] = mul l q
          bp:f32[4096] = add bn bo
          bq:f32[4096] = mul n s
          br:f32[4096] = add bp bq
          bs:f32[4096] = mul h u
          bt:f32[4096] = add br bs
          bu:f32[4096] = mul j u
          bv:f32[4096] = mul l s
          bw:f32[4096] = sub bu bv
          bx:f32[4096] = mul n q
          by:f32[4096] = add bw bx
          bz:f32[4096] = mul h w
          ca:f32[4096] = add by bz
          cb:f32[4096,1] = broadcast_in_dim[
            broadcast_dimensions=(0,)
            shape=(4096, 1)
          ] be
          cc:f32[4096,1] = broadcast_in_dim[
            broadcast_dimensions=(0,)
            shape=(4096, 1)
          ] bl
          cd:f32[4096,1] = broadcast_in_dim[
            broadcast_dimensions=(0,)
            shape=(4096, 1)
          ] bt
          ce:f32[4096,1] = broadcast_in_dim[
            broadcast_dimensions=(0,)
            shape=(4096, 1)
          ] ca
          cf:f32[4096,4] = concatenate[dimension=1] cb cc cd ce
        in (cf,) }
    ] a b
  in (c,) }

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (5189178) to head (ae9488c).
Report is 3 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master       #23      +/-   ##
===========================================
+ Coverage   99.80%   100.00%   +0.20%     
===========================================
  Files          13        13              
  Lines         500       468      -32     
===========================================
- Hits          499       468      -31     
+ Misses          1         0       -1     

Flag	Coverage Δ
unittests	100.00% <100.00%> (+0.20%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[brentyi]

Runtime changes for optimizer in EgoAllo:

• JIT: 90.1 seconds => 49.8 seconds
• Optimization: 0.178 => 0.155 seconds

On RTX 6000 Ada.