Searched refs:dimension_to_sort (Results 1 – 7 of 7) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/ |
| H A D | sort_util.cc | 155 const IrArray::Index& tiled_keys_index, int64_t dimension_to_sort, in EmitTiledCompareLoop() argument
175 b->CreateShl(tiled_keys_index[dimension_to_sort], value_one); in EmitTiledCompareLoop()
204 keys_multi_index[dimension_to_sort] = index; in EmitTiledCompareLoop()
259 b->CreateMul(tiled_keys_index[dimension_to_sort], in EmitTiledCompareLoop()
293 keys_multi_index[dimension_to_sort] = index; in EmitTiledCompareLoop()
319 int64_t dimension_to_sort, const std::vector<IrArray>& values_arrays, in EmitSortInPlace() argument
335 int64_t dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); in EmitSortInPlace()
340 if (dimension != dimension_to_sort) { in EmitSortInPlace()
346 iteration_order_to_logical_order[dim] = dimension_to_sort; in EmitSortInPlace()
390 keys_index, dimension_to_sort, dimension_to_sort_bound, xor_masks, in EmitSortInPlace()
[all …]
|
| H A D | sort_util.h | 37 int64_t dimension_to_sort, const std::vector<IrArray>& values_arrays,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | ir_emitter.h | 205 void EmitCompareLoop(int64_t dimension_to_sort,
|
| H A D | ir_emitter_unnested.cc | 2563 int64_t dimension_to_sort = sort_op.getDimension(); in EmitSort() local
2590 uint64_t dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); in EmitSort()
2629 standard_iteration_shape.set_dimensions(dimension_to_sort, in EmitSort()
2645 iteration_shape.set_dimensions(dimension_to_sort, num_iterations_in_sort_dim); in EmitSort()
2708 dimension_to_sort, values_arrays, IrName(op_name), xor_masks, &b_, in EmitSort()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | hlo_creation_utils.h | 239 int64_t dimension_to_sort, bool is_stable, HloComputation::Builder* builder,
|
| H A D | hlo_creation_utils.cc | 532 int64_t dimension_to_sort, bool is_stable, HloComputation::Builder* builder, in MakeSortHlo() argument
553 sort_shape, dimension_to_sort, operands, compare_computation, is_stable)); in MakeSortHlo()
|
| H A D | algebraic_simplifier.cc | 5988 int64_t dimension_to_sort = sort->dimensions(0); in HandleSort() local
5990 operand->shape().dimensions(dimension_to_sort) <= 1) { in HandleSort()
|