Меня интересует, как LLVM IR запустил программу Tensorflow с флагом --xla_dump_ir_to, в частности, я хотел бы изучить, как XLA обрабатывает распараллеливание.
У меня есть тестовая программа, работающая на машине с 56 процессорами, и в результирующем файле .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*))
Глядя на функцию ParallelForkJoin в исходном коде XLA, есть входной параметр num_partitions, который здесь установлен на 56.
Теперь я хотел бы получить аналогичный результат, который использует только некоторые из доступных процессоров, поэтому я добавил следующие строки в свою программу на 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 извлекает параллелизм?