Would input batch size affect auto clustering and op fusion?

[Simon King]

Hi XLA team and friends,

I built a model with tf.function and enabled XLA JIT. The model has input shape (None, 16) and all ops' the data type is int32. Basically, the model performs mathmatical operations with ops: tf.math.add, tf.math.mul, tf.math.bitwise_xor, tf.gather, tf.gather

I'd like to understand whether the input batch size will affect auto-clustering and fusion during inference time. First, I give input with batch size 10 and dump the XLA programs by setting XLA_FLAGS=--xla_dump_to. Second, I give input with a large batch size of 3162277. During execution, tensorflow gives the warning "Allocation of 3440557376 exceeds 10% of free system memory." With the same XLA_FLAGS, I dump and get another XLA programs.

By using diff command to compare the *.after_optimizations.txt files of two batch sizes, I find there are lots of difference. Basically, when with the small batch size, each fusion op seems to contain more tensorflow ops. Therefore, *.after_optimizations.txt file of the large batch size has more ops than that of small batch size. I have attached the two *.after_optimizations.txt files.

Would input batch size affect auto clustering and op fusion? If so, how does the large batch size affect auto clustering and op fusion?

[George Karpenkov]

hi Simon,

it seems that it's too late to look at HLO computations, since there's difference at clustering time already. Could you use TF_CPP_VMODULE to dump output from mark_for_compilation_pass and then use TF_DUMP_GRAPH_PREFIX to dump the clusters formed? I think you would need to look at the choices the clustering heuristic is doing, which is notoriously hard to debug.

George

--
You received this message because you are subscribed to the Google Groups "XLA development" group.
To unsubscribe from this group and stop receiving emails from it, send an email to xla-dev+u...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/xla-dev/b28ca5f4-f802-4207-b755-b0c4c3f1b02an%40googlegroups.com.