Searched refs:values_arrays (Results 1 – 3 of 3) sorted by relevance
/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/ |
D | sort_util.cc | 295 int64 dimension_to_sort, const std::vector<IrArray>& values_arrays, in EmitSortInPlace() argument 309 const Shape& keys_shape = values_arrays[0].GetShape(); in EmitSortInPlace() 328 std::vector<llvm::Value*> param_shmem_buffers(values_arrays.size(), nullptr); in EmitSortInPlace() 331 for (int64 i = 0; i < values_arrays.size(); ++i) { in EmitSortInPlace() 334 values_arrays[i].GetShape().element_type(), module), in EmitSortInPlace() 362 IrArray::Index keys_index(keys_multi_index, values_arrays[0].GetShape(), in EmitSortInPlace() 366 values_arrays, param_shmem_buffers, tile_size, emit_compare_callback, in EmitSortInPlace() 372 values_arrays[operand].GetShape(), in EmitSortInPlace() 374 return values_arrays[operand].EmitArrayElementAddress(keys_index, b); in EmitSortInPlace() 380 values_arrays[operand].GetShape(), in EmitSortInPlace() [all …]
|
D | sort_util.h | 38 int64 dimension_to_sort, const std::vector<IrArray>& values_arrays,
|
/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
D | ir_emitter_unnested.cc | 2999 std::vector<IrArray> values_arrays; in EmitSortFromMlir() local 3000 values_arrays.reserve(context.operand_shapes.size()); in EmitSortFromMlir() 3002 values_arrays.push_back(ir_arrays[i]); in EmitSortFromMlir() 3008 dimension_to_sort, values_arrays, IrName(context.name), xor_masks, &b_, in EmitSortFromMlir()
|