我对通过使用--xla_dump_ir_to
标志启动Tensorflow程序而获得的LLVM IR感兴趣,特别是我想探索XLA如何处理并行化。
我有一个测试程序在具有56个CPU的计算机上运行,并且在生成的.ll
文件中,按预期方式显示了并行化功能:
call void @__xla_cpu_runtime_ParallelForkJoin(i8* %6, i8* %run_options, i8** null, i8** %buffer_table, i64* %prof_counters, i32 56, i64* getelementptr inbounds ([224 x i64], [224 x i64]* @parallel_convolution_parallel_dimension_partitions, i32 0, i32 0), i32 2, i8* bitcast (void (i8*, i8*, i8**, i8**, i64*, i64*)* @parallel_convolution to i8*))
查看XLA源代码中的ParallelForkJoin
函数,这里有一个输入参数num_partitions
设置为56。
现在,我想获得仅使用某些可用CPU的类似结果,因此我在python程序中添加了以下几行(如Tensorflow指南https://www.tensorflow.org/guide/performance/overview#optimizing_for_cpu和SO上其他地方的建议):
config = tf.ConfigProto()
config.intra_op_parallelism_threads = 5
config.inter_op_parallelism_threads = 5
tf.Session(config=config)
但是生成的.ll
文件完全相同!
那是为什么?我本以为,指定不同数量的并行线程会导致参数5(在这种情况下)而不是56。
一般而言,是否有关于XLA如何提取并行性的文档?