xla

cdc8f584 - improve collective matmul perf