add batch_norm op with test and benchmark

[yanghailong-git]

Summary

Implemented a 2D batch normalization Triton operator, successfully ran the corresponding tests and benchmarks, and visualized the performance tests for speed and memory.

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

the visualization of performance:

[yundai424]

looks like from the benchmark result triton impl is slower than HF original one? 👀

[yanghailong-git]

looks like from the benchmark result triton impl is slower than HF original one? 👀

It seems so. The memory usage is about the same, but the speed is a bit slower. Do you have any optimization or improvement methods?