Add Mooncake rule for propagate e_mul_xj fast path and added tests

[Parvm1102]

I have added the propagate e_mul_xj fast path. This pr depends on #677 so that should be merged first.

Benchmark Results

The rules matched for the following layers, the improvements and the benchmarks are listed below, all layers pass the gradient correctness benchmarks:

GCNConv

[n=512]
  grad-check PASS

  Zygote             Zyg     1.460 ms /   2501 alloc
  Mooncake + RULE     MC     2.882 ms /    473 alloc   (MC+rule / Zyg =  1.97x)
  Mooncake - rule     MC     9.504 ms /    665 alloc   (no-rule / Zyg =  6.51x)

  >>> RULE SPEEDUP (no-rule / +rule) =  3.30x

[n=2048]
  grad-check PASS

  Zygote             Zyg     7.601 ms /   2501 alloc
  Mooncake + RULE     MC    15.561 ms /    473 alloc   (MC+rule / Zyg =  2.05x)
  Mooncake - rule     MC    39.605 ms /    665 alloc   (no-rule / Zyg =  5.21x)

  >>> RULE SPEEDUP (no-rule / +rule) =  2.55x

[n=8192]
  grad-check PASS

  Zygote             Zyg   125.052 ms /   2501 alloc
  Mooncake + RULE     MC   132.897 ms /    473 alloc   (MC+rule / Zyg =  1.06x)
  Mooncake - rule     MC   467.868 ms /    665 alloc   (no-rule / Zyg =  3.74x)

  >>> RULE SPEEDUP (no-rule / +rule) =  3.52x

SGConv (k=2)

[n=512]
  grad-check PASS

  Zygote             Zyg     2.791 ms /   4150 alloc
  Mooncake + RULE     MC     4.762 ms /    560 alloc   (MC+rule / Zyg =  1.71x)
  Mooncake - rule     MC    18.177 ms /    946 alloc   (no-rule / Zyg =  6.51x)

  >>> RULE SPEEDUP (no-rule / +rule) =  3.82x

[n=2048]
  grad-check PASS

  Zygote             Zyg    10.970 ms /   4150 alloc
  Mooncake + RULE     MC    18.078 ms /    560 alloc   (MC+rule / Zyg =  1.65x)
  Mooncake - rule     MC    75.055 ms /    946 alloc   (no-rule / Zyg =  6.84x)

  >>> RULE SPEEDUP (no-rule / +rule) =  4.15x

[n=8192]
  grad-check PASS

  Zygote             Zyg   233.485 ms /   4150 alloc
  Mooncake + RULE     MC   206.554 ms /    560 alloc   (MC+rule / Zyg =  0.88x)
  Mooncake - rule     MC   944.392 ms /    946 alloc   (no-rule / Zyg =  4.04x)

  >>> RULE SPEEDUP (no-rule / +rule) =  4.57x

TAGConv (k=2)

[n=512]
  grad-check PASS

  Zygote             Zyg     3.412 ms /   4195 alloc
  Mooncake + RULE     MC     5.849 ms /    593 alloc   (MC+rule / Zyg =  1.71x)
  Mooncake - rule     MC    20.454 ms /    979 alloc   (no-rule / Zyg =  5.99x)

  >>> RULE SPEEDUP (no-rule / +rule) =  3.50x

[n=2048]
  grad-check PASS

  Zygote             Zyg    12.748 ms /   4195 alloc
  Mooncake + RULE     MC    29.878 ms /    593 alloc   (MC+rule / Zyg =  2.34x)
  Mooncake - rule     MC    75.978 ms /    979 alloc   (no-rule / Zyg =  5.96x)

  >>> RULE SPEEDUP (no-rule / +rule) =  2.54x

[n=8192]
  grad-check PASS

  Zygote             Zyg   198.058 ms /   4195 alloc
  Mooncake + RULE     MC   267.873 ms /    593 alloc   (MC+rule / Zyg =  1.35x)
  Mooncake - rule     MC   957.197 ms /    979 alloc   (no-rule / Zyg =  4.83x)

  >>> RULE SPEEDUP (no-rule / +rule) =  3.57x

I have also added an extra test for this change.

[Parvm1102]

@CarloLucibello please review this after #677

[CarloLucibello]

this seems fine, just needs rebase

[wsmoses]

Mind also adding an enzyme rule too (in addition to those in the other PRs)

[CarloLucibello]

Mind also adding an enzyme rule too (in addition to those in the other PRs)

we will tackle this after mooncake is done