Home
last modified time | relevance | path

Searched refs:IrArray (Results 1 – 25 of 36) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/
Dloop_emitter.cc47 const IrArray& target_array, llvm::IRBuilder<>* b) in LoopEmitter()
48 : body_emitter_([=](const llvm_ir::IrArray::Index array_index) -> Status { in LoopEmitter()
60 const std::vector<IrArray>& target_arrays, llvm::IRBuilder<>* b) { in MakeBodyEmitterForMultiOutput()
61 return [=](const llvm_ir::IrArray::Index array_index) { in MakeBodyEmitterForMultiOutput()
79 absl::Span<const IrArray> target_arrays, in LoopEmitter()
83 std::vector<IrArray>(target_arrays.begin(), target_arrays.end()), b)), in LoopEmitter()
88 for (const IrArray& array : target_arrays) { in LoopEmitter()
95 IrArray::Index LoopEmitter::EmitStaticIndex(ForLoopNest* loop_nest, in EmitStaticIndex()
110 return IrArray::Index(array_multi_index, shape_, index_type); in EmitStaticIndex()
113 IrArray::Index LoopEmitter::EmitDynamicIndex(ForLoopNest* loop_nest, in EmitDynamicIndex()
[all …]
Dloop_emitter.h37 std::function<StatusOr<llvm::Value*>(const IrArray::Index& index)>;
42 using BodyEmitter = std::function<Status(const IrArray::Index& index)>;
55 const IrArray& target_array, llvm::IRBuilder<>* b);
63 absl::Span<const IrArray> target_arrays, llvm::IRBuilder<>* b);
73 std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock() { in EmitIndexAndSetExitBasicBlock()
78 virtual std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock(
104 IrArray::Index EmitStaticIndex(ForLoopNest* loop_nest,
106 IrArray::Index EmitDynamicIndex(ForLoopNest* loop_nest,
Dir_array.cc37 IrArray::Index::Index(absl::Span<llvm::Value* const> multidim, in Index()
45 void IrArray::Index::Delinearize(std::vector<llvm::Value*>* multidim, in Delinearize()
75 void IrArray::Index::Delinearize(std::vector<llvm::Value*>* multidim, in Delinearize()
101 IrArray::Index::Index(llvm::Value* linear, const Shape& shape, in Index()
115 IrArray::Index::Index(llvm::Value* linear, const Shape& shape, in Index()
130 IrArray::Index::Index(absl::Span<llvm::Value* const> multidim, in Index()
136 IrArray::Index::Index(absl::Span<llvm::Value* const> multidim, in Index()
153 IrArray::IrArray(llvm::Value* base_ptr, Shape shape) in IrArray() function in xla::llvm_ir::IrArray
174 bool IrArray::Index::LinearValidOnShape(const Shape& a) const { in LinearValidOnShape()
182 IrArray::Index IrArray::Index::SourceIndexOfReshape( in SourceIndexOfReshape()
[all …]
Ddynamic_update_slice_util.cc98 const IrArray& output_array, const gpu::LaunchDimensions* launch_dimensions, in EmitDynamicUpdateSliceInPlaceImpl()
130 auto loop_body_emitter = [&](const IrArray::Index& update_index) -> Status { in EmitDynamicUpdateSliceInPlaceImpl()
144 IrArray::Index output_index(output_multi_index, output_shape, in EmitDynamicUpdateSliceInPlaceImpl()
160 Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays, in EmitDynamicUpdateSliceInPlace()
161 const IrArray& output_array, in EmitDynamicUpdateSliceInPlace()
168 IrArray update_array = operand_arrays[1]; in EmitDynamicUpdateSliceInPlace()
169 IrArray start_indices_array = operand_arrays[2]; in EmitDynamicUpdateSliceInPlace()
175 IrArray::Index(b->getInt64Ty()), b); in EmitDynamicUpdateSliceInPlace()
177 ElementGenerator update_array_generator = [&](const IrArray::Index& index) { in EmitDynamicUpdateSliceInPlace()
192 const HloComputation* fusion, const IrArray& fusion_output_array, in EmitFusedDynamicUpdateSliceInPlaceImpl()
[all …]
Dir_array.h44 class IrArray {
211 IrArray() : base_ptr_(nullptr) {} in IrArray() function
215 IrArray(llvm::Value* base_ptr, Shape shape);
218 IrArray(IrArray&& other) = default;
219 IrArray(const IrArray& other) = default;
220 IrArray& operator=(IrArray&& other) = default;
221 IrArray& operator=(const IrArray& other) = default;
266 IrArray CastToShape(const Shape& new_shape, llvm::IRBuilder<>* b) const;
Dtuple_ops.h62 void EmitTupleSelect(const IrArray& select, const IrArray& pred,
68 void EmitTuple(const IrArray& tuple, absl::Span<llvm::Value* const> operands,
79 void EmitTuple(const IrArray& tuple, absl::Span<const IrArray> buffers,
Ddynamic_update_slice_util.h32 std::function<std::vector<llvm_ir::IrArray>()>;
66 Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays,
67 const IrArray& output_array,
76 const IrArray& fusion_output_array,
83 const HloComputation* fusion, const IrArray& fusion_output_array,
Dir_array_test.cc33 EXPECT_TRUE(IrArray::Index::ShapeIsCompatible(a, b)); in TEST()
34 EXPECT_TRUE(IrArray::Index::ShapeIsCompatible(a, c)); in TEST()
35 EXPECT_FALSE(IrArray::Index::ShapeIsCompatible(a, d)); in TEST()
36 EXPECT_FALSE(IrArray::Index::ShapeIsCompatible(a, e)); in TEST()
37 EXPECT_FALSE(IrArray::Index::ShapeIsCompatible(a, f)); in TEST()
Dtuple_ops.cc36 void EmitTupleSelect(const IrArray& select, const IrArray& pred, in EmitTupleSelect()
61 void EmitTuple(const IrArray& tuple, absl::Span<llvm::Value* const> operands, in EmitTuple()
74 void EmitTuple(const IrArray& tuple, absl::Span<const IrArray> buffers, in EmitTuple()
80 [](const llvm_ir::IrArray& buffer) { return buffer.GetBasePointer(); }); in EmitTuple()
Dsort_util.cc146 const IrArray::Index& tiled_keys_index, int64 dimension_to_sort, in EmitTiledCompareLoop()
148 const std::vector<IrArray>& params, in EmitTiledCompareLoop()
195 IrArray::Index keys_index(keys_multi_index, params[i].GetShape(), in EmitTiledCompareLoop()
274 IrArray::Index keys_index(keys_multi_index, params[i].GetShape(), in EmitTiledCompareLoop()
295 int64 dimension_to_sort, const std::vector<IrArray>& values_arrays, in EmitSortInPlace()
342 [&](const IrArray::Index& tiles_index) -> Status { in EmitSortInPlace()
362 IrArray::Index keys_index(keys_multi_index, values_arrays[0].GetShape(), in EmitSortInPlace()
371 IrArray::Index keys_index(keys_multi_index, in EmitSortInPlace()
379 IrArray::Index keys_index(keys_multi_index, in EmitSortInPlace()
Dfused_ir_emitter.cc41 using llvm_ir::IrArray;
45 [=](const IrArray::Index& index) -> StatusOr<llvm::Value*> { in DefaultAction()
85 indexed_generators_[constant] = [=](const IrArray::Index& index) { in HandleConstant()
104 return IrArray(shape_constant, constant->shape()) in HandleConstant()
131 [=](const IrArray::Index& index) -> StatusOr<llvm::Value*> { in HandleTuple()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.h135 const llvm_ir::IrArray::Index& index, llvm::Value* y_loc,
142 const ThreadIdInfo& thread_id_info, const llvm_ir::IrArray::Index& index,
366 absl::Span<const llvm_ir::IrArray> result_ir_arrays,
367 const llvm_ir::IrArray::Index& index, bool use_linear_index,
458 absl::Span<const llvm_ir::IrArray> ir_arrays,
459 const llvm_ir::IrArray::Index& index);
469 const llvm_ir::IrArray& output,
485 llvm_ir::IrArray output;
504 absl::Span<const llvm_ir::IrArray> operand_arrays,
505 absl::Span<const llvm_ir::IrArray> output_arrays,
[all …]
Dir_emitter.h73 std::function<std::vector<llvm_ir::IrArray>()>;
125 llvm_ir::IrArray GetIrArray(const HloInstruction& inst,
138 std::vector<llvm_ir::IrArray> ConstructIrArrayForOutputs(
207 const llvm_ir::IrArray::Index& keys_index,
208 const llvm_ir::IrArray::Index& compare_keys_index,
209 const llvm_ir::IrArray& keys_array);
Dir_emitter_unnested.cc138 using llvm_ir::IrArray;
820 std::vector<llvm_ir::IrArray> ir_arrays; in EmitPadToStaticFromMlir()
826 const llvm_ir::IrArray source_array = ir_arrays[0]; in EmitPadToStaticFromMlir()
827 const llvm_ir::IrArray output_array = ir_arrays[1]; in EmitPadToStaticFromMlir()
829 absl::Span<const llvm_ir::IrArray>(ir_arrays).subspan(2); in EmitPadToStaticFromMlir()
907 [&](const llvm_ir::IrArray::Index& array_index) -> Status { in EmitPadToStaticFromMlir()
915 llvm_ir::IrArray::Index dyn_index(linearIndex, input_shape, in EmitPadToStaticFromMlir()
947 std::vector<llvm_ir::IrArray> ir_arrays; in EmitSliceToDynamicFromMlir()
973 const llvm_ir::IrArray data_array = ir_arrays.back(); in EmitSliceToDynamicFromMlir()
1028 [&](const llvm_ir::IrArray::Index& array_index) -> Status { in EmitSliceToDynamicFromMlir()
[all …]
Dparallel_loop_emitter.h41 const llvm_ir::IrArray& target_array,
51 absl::Span<const llvm_ir::IrArray> target_arrays,
59 std::vector<llvm_ir::IrArray::Index> EmitIndexAndSetExitBasicBlock(
Dparallel_loop_emitter.cc44 absl::Span<const llvm_ir::IrArray> target_arrays, in ParallelLoopEmitter()
53 const llvm_ir::IrArray& target_array, in ParallelLoopEmitter()
60 std::vector<llvm_ir::IrArray::Index>
77 std::vector<llvm_ir::IrArray::Index> array_indices; in EmitIndexAndSetExitBasicBlock()
180 for (const llvm_ir::IrArray::Index& array_index : in EmitLoop()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Ddot_op_emitter.cc125 const llvm_ir::IrArray& target_array,
126 const llvm_ir::IrArray& lhs_array,
127 const llvm_ir::IrArray& rhs_array,
128 const llvm_ir::IrArray* addend_array,
223 const llvm_ir::IrArray& target_array_;
224 const llvm_ir::IrArray& lhs_array_;
225 const llvm_ir::IrArray& rhs_array_;
226 const llvm_ir::IrArray* addend_array_;
236 DotInfo dot_info, string dot_hlo_name, const llvm_ir::IrArray& target_array, in DotOpEmitter()
237 const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, in DotOpEmitter()
[all …]
Ddot_op_emitter.h62 const llvm_ir::IrArray& target_array,
63 const llvm_ir::IrArray& lhs_array,
64 const llvm_ir::IrArray& rhs_array,
65 const llvm_ir::IrArray* addend_array,
Dparallel_loop_emitter.cc27 const llvm_ir::IrArray& target_array, in ParallelLoopEmitter()
32 std::vector<llvm_ir::IrArray::Index>
80 llvm_ir::IrArray::Index array_index(array_multi_index, shape_, index_type); in EmitIndexAndSetExitBasicBlock()
Dir_emitter.cc120 llvm_ir::IrArray root_value = GetIrArrayFor(computation->root_instruction()); in EmitThreadLocalFunctionEpilogue()
438 llvm_ir::EmitTuple(llvm_ir::IrArray(data_address, data_shape), in HandleInfeed()
704 [this, init_value](const llvm_ir::IrArray::Index& target_index) { in HandleSelectAndScatter()
711 const llvm_ir::IrArray::Index source_index = in HandleSelectAndScatter()
734 const llvm_ir::IrArray::Index window_index = window_loops.AddLoopsForShape( in HandleSelectAndScatter()
768 [&](const llvm_ir::IrArray::Index& operand_index) { in HandleSelectAndScatter()
775 llvm_ir::IrArray operand_array(GetIrArrayFor(operand)); in HandleSelectAndScatter()
776 llvm_ir::IrArray::Index operand_index( in HandleSelectAndScatter()
817 llvm_ir::IrArray source_array(GetIrArrayFor(source)); in HandleSelectAndScatter()
820 llvm_ir::IrArray output_array(GetIrArrayFor(select_and_scatter)); in HandleSelectAndScatter()
[all …]
Dparallel_loop_emitter.h55 const llvm_ir::IrArray& target_array,
63 std::vector<llvm_ir::IrArray::Index> EmitIndexAndSetExitBasicBlock(
Dir_emitter.h69 std::function<std::vector<llvm_ir::IrArray>()>;
225 llvm_ir::IrArray GetIrArrayFor(const HloInstruction* hlo);
228 std::vector<llvm_ir::IrArray> GetIrArraysForOperandsOf(
237 llvm_ir::IrArray* array) { in AddAliasingInformationToIrArray()
383 const llvm_ir::IrArray& containing_array);
399 const llvm_ir::IrArray::Index& output_index,
415 const llvm_ir::IrArray& target_array,
416 const llvm_ir::IrArray& source_array);
/external/tensorflow/tensorflow/compiler/xla/service/cpu/tests/
Dcpu_noalias_test.cc93 llvm_ir::IrArray param_x_array(param_x_val, param_shape); in TEST_F()
95 llvm_ir::IrArray::Index zero_2d({zero, zero}, param_shape, zero->getType()); in TEST_F()
104 llvm_ir::IrArray concat1_array(concat1_val, shape); in TEST_F()
106 llvm_ir::IrArray::Index zero_2d({zero, zero}, shape, zero->getType()); in TEST_F()
115 llvm_ir::IrArray concat2_array(concat2_val, shape); in TEST_F()
117 llvm_ir::IrArray::Index zero_2d({zero, zero}, shape, zero->getType()); in TEST_F()
125 llvm_ir::IrArray add_array(concat2_val, shape); in TEST_F()
127 llvm_ir::IrArray::Index zero_2d({zero, zero}, shape, zero->getType()); in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/
Delemental_ir_emitter.h206 const llvm_ir::IrArray::Index& index);
211 const llvm_ir::IrArray::Index& index);
216 const llvm_ir::IrArray::Index& target_index);
221 const llvm_ir::IrArray::Index& index);
226 const llvm_ir::IrArray::Index& index);
231 const llvm_ir::IrArray::Index& index);
236 const llvm_ir::IrArray::Index& padded_index);
241 const llvm_ir::IrArray::Index& dot_result_index);
255 const llvm_ir::IrArray::Index& index);
261 const llvm_ir::IrArray::Index& index);
[all …]
Delemental_ir_emitter.cc54 using llvm_ir::IrArray;
1718 const llvm_ir::IrArray::Index& index) { in EmitElementalSelect()
1732 const llvm_ir::IrArray::Index& index) { in EmitElementalClamp()
1755 const llvm_ir::IrArray::Index& target_index) { in EmitElementalConcatenate()
1826 llvm_ir::IrArray::Index operand_index(operand_multi_index, operand->shape(), in EmitElementalConcatenate()
1862 const llvm_ir::IrArray::Index& index) { in EmitElementalDynamicSlice()
1873 llvm_ir::IrArray::Index zero_index(index_type); in EmitElementalDynamicSlice()
1901 llvm_ir::IrArray::Index input_index(input_multi_index, input_hlo->shape(), in EmitElementalDynamicSlice()
1909 const llvm_ir::IrArray::Index& index) { in EmitElementalGather()
2001 IrArray::Index gather_index_index(gather_index_index_components, in EmitElementalGather()
[all …]

12