Refactor GELU and Sigmoid epilogue to use a common template (and add SiLu, Hardswish epilogue)

[masahi]

Since Sigmoid and GeLU epilogue functors are almost identical, I propose to refactor them. The motivation is to make it easy to add another functors.

Now Sigmoid and GeLU epilogue functors are defined as:

using LinearCombinationSigmoid = LinearCombinationGeneric<Sigmoid, ElementOutput_, Count, ElementAccumulator_,
                                                          ElementCompute_, FloatRoundStyle::round_to_nearest>;

using LinearCombinationGELU = LinearCombinationGeneric<GELU, ElementOutput_, Count, ElementAccumulator_,
                                                          ElementCompute_, FloatRoundStyle::round_to_nearest>;

I'm going to add two new functors for SiLU and HardSwish activation, which can be implemented in the same manner.

What do you think? @hwu36

[masahi]

ok I've pushed two activation using this common template.

• SiLU https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html
• HardSwish https://pytorch.org/docs/stable/generated/torch.nn.Hardswish.html

[d-k-b]

Thanks for helping clean this up! Can you add documentation with a short description and links to resources to the SiLu and HardSwish structs for future reference?

[masahi]

@d-k-b Thanks for the feedback, the doc is updated.

[hwu36]

Overall, looks good to me. I think just change FloatRoundStyle::round_to_nearest to Round everywhere.

[masahi]

Oops thanks for spotting this, fixed.

[masahi]

Not sure why Ci has failed

[hwu36]

These 4 activations look expensive to me. You can try to set kIsHeavy to true, check the generated code is different, and compare the performance to verify it.

If you setting them to true, the epilogue will not be fully unrolled which is good for the I$. Usually complex activations such as gelu need to be set to true.

[masahi]

These 4 activations look expensive to me. You can try to set kIsHeavy to true, check the generated code is different, and compare the performance to verify it.

If you setting them to true, the epilogue will not be fully unrolled which is good for the I$. Usually complex activations such as gelu need to be set to true.

My patch respects the current choice for the existing epilogues: isHeavy = true for GELU

cutlass/include/cutlass/epilogue/thread/linear_combination_gelu.h

Line 57 in 80e8fe7

ElementCompute_, Round, true>;

and whatever the default for Sigmoid which I assumed was false .

cutlass/include/cutlass/epilogue/threadblock/epilogue_base.h

Lines 73 to 76 in 80e8fe7

	// If the epilogue functor does not define `kIsHeavy` or if it is `false`, then
	// the behavior from CUTLASS 2.5 and before is retained. The epilogue is fully
	// unrolled and inlined.
	//

For the two new activations, my reasoning was false for Silu because it is just sigmoid + mul, and also false for hardswish since it is similar to relu.

[hwu36]

SGTM.

kIsHeavy was newly introduced. We benchmarked gelu, but not sigmoid. If you find sigmoid is expensive in the future, you can create a PR and flip the bool.

[masahi]

Making sigmoid and silu isHeavy = true did improve performance on efficientnetv2 slightly: 8.58 msec -> 8.26 msec where 8.58 was the number from my PR apache/tvm#9746. So I pushed that change. The change didn't help when the vectorized variant of sigmoid is used.

[hwu36]

The change didn't help when the vectorized variant of sigmoid is used.

Vectorized one has less instruction and less burden to I$.

[hwu36]

Making sigmoid and silu isHeavy = true did improve performance on efficientnetv2 slightly: 8.58 msec -> 8.26 msec where 8.58 was the number from my PR apache/tvm#9746.

You can update your TVM PR now. :)