[Performance Drop] g.update_all() slows much down after commit acb4eb7

[Rhett-Ying]

🐛 Bug

Benchmark test bench_builtin_update_all_csc.py slows much down after commit acb4eb7.

dgl/benchmarks/benchmarks/api/bench_builtin_update_all_csc.py

Lines 12 to 18 in ae8cbde

	@utils.benchmark("time", timeout=600)
	@utils.parametrize("graph_name", ["ogbn-arxiv", "reddit", "ogbn-proteins"])
	@utils.parametrize("format", ["csc"])
	@utils.parametrize("feat_size", [4, 32, 256])
	@utils.parametrize("msg_type", ["copy_u", "u_mul_e"])
	@utils.parametrize("reduce_type", ["sum", "mean", "max"])
	def track_time(graph_name, format, feat_size, msg_type, reduce_type):

This perf issue happens for u_mul_e only, copy_u works well. See more details here.

I have reproduced on my local instance r6i.metal. below are details.
Results on commit acb4eb7

[100.00%] ··· =============== =========================== ============================
-- format / feat_size / msg_type / reduce_type
--------------- --------------------------------------------------------
graph_name csc / 256 / u_mul_e / sum csc / 256 / u_mul_e / mean
=============== =========================== ============================
reddit 11.187968121399171 11.947775951796212
ogbn-proteins 3.828348464198643 4.048665992595488
=============== =========================== ============================

Results on commit 29e6661

[100.00%] ··· =============== =========================== ============================
-- format / feat_size / msg_type / reduce_type
--------------- --------------------------------------------------------
graph_name csc / 256 / u_mul_e / sum csc / 256 / u_mul_e / mean
=============== =========================== ============================
reddit 3.305417302099522 3.3533205801039
ogbn-proteins 1.1776576915057375 1.1793251131020952
=============== =========================== ============================

To Reproduce

Steps to reproduce the behavior:

1. checkout to target commit.
2. build DGL. For example: bash script/create_dev_conda_env.sh -c -p 3.10 -s -t 1.13.1 + conda activate dgl-dev-cpu + bash script/build_dgl.sh -c.
3. cd benchmarks
4. DGL_HOME=/home/ubuntu/workspace/dgl_bench DGL_LIBRARY_PATH=$DGL_HOME/build PYTHONPATH=$DGL_HOME/python:$PYTHONPATH asv run -n -e --python=same --verbose --bench api.bench_builtin_update_all_csc

Expected behavior

Environment

• DGL Version (e.g., 1.0):
• Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
• OS (e.g., Linux):
• How you installed DGL (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version (if applicable):
• GPU models and configuration (e.g. V100):
• Any other relevant information:

Additional context

[Rhett-Ying]

@itaraban could you help look into this issue as you're the owner of #5497? This is blocking us cutting a new release of DGL.