Introduce @skip_rewrite_func, @skip_rewrite_type to extend should_rewrite_call

[mofeing]

Performance comparison

I'm using Muscle.binary_einsum to compare compile time. It is semantically equal to stablehlo.dot_general. What I found is that the compilation time disappears on the second time a function is traced. Notably I get a x5-30 comptime speedup on second tracing. Most probably this recompilation that disappears might be because Muscle.binary_einsum is a type-unstable function and AFAIK that is a dynamic call that Reactant always rewrites. In that case, this PR successfully avoids the rewrites on those cases.

Unfortunately, first trace comp time doesn't get lower and the bottleneck is probably in other part.

julia> a = Tensor(ConcreteRArray(rand(3, 4)), [Index(:i), Index(:j)])
3×4 Tensor{Float64, 2, ConcretePJRTArray{Float64, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}:
 0.00641709  0.247586  0.376454  0.75816
 0.896955    0.577077  0.947384  0.97931
 0.60413     0.567416  0.804769  0.489075

julia> b = Tensor(ConcreteRArray(rand(3, 4)), [Index(:i), Index(:j)])
3×4 Tensor{Float64, 2, ConcretePJRTArray{Float64, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}}:
 0.196665   0.198938  0.421027  0.0856663
 0.0510834  0.430934  0.264092  0.625037
 0.488447   0.624371  0.28303   0.493439

julia> g(x,y) = binary_einsum(x,y)
g (generic function with 1 method)

Before

julia> @time @code_hlo optimize=false g(a, b);
 14.256694 seconds (52.07 M allocations: 2.618 GiB, 5.30% gc time, 99.76% compilation time: <1% of which was recompilation)

julia> @time @code_hlo optimize=false g(a, b);
  0.093615 seconds (5.86 k allocations: 462.086 KiB, 95.72% compilation time)

After

julia> @time @code_hlo optimize=false g(a, b);
 13.495342 seconds (52.50 M allocations: 2.642 GiB, 6.19% gc time, 99.80% compilation time: <1% of which was recompilation)

julia> @time @code_hlo optimize=false g(a, b);
  0.003218 seconds (3.22 k allocations: 359.414 KiB)

Open questions

• If the top-function in the traced function (e.g. f in @compile f(x)) is marked, it won't be detected and it will rewrite it. Before starting the rewrite, we should check if the top-function must be rewritten. Here the proof that it's not working:

julia> @time @code_hlo optimize=false binary_einsum(a, b);
  1.452055 seconds (2.54 M allocations: 133.979 MiB, 1.04% gc time, 99.55% compilation time)

julia> @time @code_hlo optimize=false binary_einsum(a, b);
  0.020816 seconds (3.66 k allocations: 377.211 KiB, 84.59% compilation time)

• Is there any way we can mark only certain methods to skip the rewrite (unlike the whole function)? i.e. extend these changes to should_rewrite_invoke.
  • The only idea to implement this that comes to my mind is using another Core.MethodTable that it's only used to check methods that skip rewritting.

[wsmoses]

can you add a test to confirm?

[mofeing]

can you add a test to confirm?

done, i'm not sure how to test @skip_rewrite_type but @skip_rewrite_func is tested now

[mofeing]

Failed tests seem unrelated to the PR.