Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/
Dloop_emitter.cc39 const IrArray& target_array, llvm::IRBuilder<>* b) in LoopEmitter()
40 : body_emitter_([=](const llvm_ir::IrArray::Index array_index) -> Status { in LoopEmitter()
52 const std::vector<IrArray>& target_arrays, llvm::IRBuilder<>* b) { in MakeBodyEmitterForMultiOutputFusion()
53 return [=](const llvm_ir::IrArray::Index array_index) { in MakeBodyEmitterForMultiOutputFusion()
71 absl::Span<const IrArray> target_arrays, in LoopEmitter()
75 std::vector<IrArray>(target_arrays.begin(), target_arrays.end()), b)), in LoopEmitter()
80 for (const IrArray& array : target_arrays) { in LoopEmitter()
87 std::vector<IrArray::Index> LoopEmitter::EmitIndexAndSetExitBasicBlock( in EmitIndexAndSetExitBasicBlock()
93 return {IrArray::Index(index_type)}; in EmitIndexAndSetExitBasicBlock()
110 IrArray::Index array_index(array_multi_index, shape_, index_type); in EmitIndexAndSetExitBasicBlock()
[all …]
Dir_array.cc32 IrArray::Index::Index(absl::Span<llvm::Value* const> multidim, in Index()
40 void IrArray::Index::Delinearize(std::vector<llvm::Value*>* multidim, in Delinearize()
70 IrArray::Index::Index(llvm::Value* linear, const Shape& shape, in Index()
84 IrArray::Index::Index(absl::Span<llvm::Value* const> multidim, in Index()
101 IrArray::IrArray(llvm::Value* base_ptr, Shape shape) in IrArray() function in xla::llvm_ir::IrArray
122 bool IrArray::Index::LinearValidOnShape(const Shape& a) const { in LinearValidOnShape()
130 IrArray::Index IrArray::Index::SourceIndexOfReshape( in SourceIndexOfReshape()
180 IrArray::Index IrArray::Index::SourceIndexOfSlice( in SourceIndexOfSlice()
201 IrArray::Index IrArray::Index::SourceIndexOfTranspose( in SourceIndexOfTranspose()
217 IrArray::Index IrArray::Index::SourceIndexOfBitcast( in SourceIndexOfBitcast()
[all …]
Ddynamic_update_slice_util.cc44 const IrArray& output_array, const gpu::LaunchDimensions* launch_dimensions, in EmitDynamicUpdateSliceInPlaceImpl()
76 auto loop_body_emitter = [&](const IrArray::Index& update_index) -> Status { in EmitDynamicUpdateSliceInPlaceImpl()
90 IrArray::Index output_index(output_multi_index, output_shape, in EmitDynamicUpdateSliceInPlaceImpl()
106 Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays, in EmitDynamicUpdateSliceInPlace()
107 const IrArray& output_array, in EmitDynamicUpdateSliceInPlace()
114 IrArray update_array = operand_arrays[1]; in EmitDynamicUpdateSliceInPlace()
115 IrArray start_indices_array = operand_arrays[2]; in EmitDynamicUpdateSliceInPlace()
121 IrArray::Index(b->getInt64Ty()), b); in EmitDynamicUpdateSliceInPlace()
123 ElementGenerator update_array_generator = [&](const IrArray::Index& index) { in EmitDynamicUpdateSliceInPlace()
140 const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter, in EmitFusedDynamicUpdateSliceInPlaceImpl()
[all …]
Dir_array.h43 class IrArray {
204 IrArray() : base_ptr_(nullptr) {} in IrArray() function
208 IrArray(llvm::Value* base_ptr, Shape shape);
211 IrArray(IrArray&& other) = default;
212 IrArray(const IrArray& other) = default;
213 IrArray& operator=(IrArray&& other) = default;
214 IrArray& operator=(const IrArray& other) = default;
259 IrArray CastToShape(const Shape& new_shape, llvm::IRBuilder<>* b) const;
Dkernel_tiling.cc58 IrArray::Index GetReshapedIndex(const IrArray::Index& index, const Shape& shape, in GetReshapedIndex()
75 return IrArray::Index(linear_index, reshaped_shape, b); in GetReshapedIndex()
149 IrArray::Index KernelMappingScheme::GetUnnormalizedIndex( in GetUnnormalizedIndex()
150 const IrArray::Index& normalized_shape_index, in GetUnnormalizedIndex()
159 IrArray::Index KernelMappingScheme::EmitBlockIndex(llvm::Type* index_ty) { in EmitBlockIndex()
166 return IrArray::Index(linear_block_id, in EmitBlockIndex()
172 IrArray::Index KernelMappingScheme::GetTileIndexForBlockOrigin( in GetTileIndexForBlockOrigin()
173 const IrArray::Index& block_index) { in GetTileIndexForBlockOrigin()
183 return IrArray::Index(multidim, block_index[0]->getType()); in GetTileIndexForBlockOrigin()
186 IrArray::Index KernelMappingScheme::GetElementIndexForTileOrigin( in GetElementIndexForTileOrigin()
[all …]
Dloop_emitter.h36 std::function<StatusOr<llvm::Value*>(const IrArray::Index& index)>;
41 using BodyEmitter = std::function<Status(const IrArray::Index& index)>;
48 const IrArray& target_array, llvm::IRBuilder<>* b);
56 absl::Span<const IrArray> target_arrays, llvm::IRBuilder<>* b);
66 std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock() { in EmitIndexAndSetExitBasicBlock()
70 virtual std::vector<IrArray::Index> EmitIndexAndSetExitBasicBlock(
Dtuple_ops.h62 void EmitTupleSelect(const IrArray& select, const IrArray& pred,
68 void EmitTuple(const IrArray& tuple, absl::Span<llvm::Value* const> operands,
73 void EmitTuple(const IrArray& tuple, absl::Span<const IrArray> buffers,
Ddynamic_update_slice_util.h31 std::function<std::vector<llvm_ir::IrArray>()>;
69 Status EmitDynamicUpdateSliceInPlace(absl::Span<const IrArray> operand_arrays,
70 const IrArray& output_array,
81 const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter,
89 const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter,
Dkernel_tiling.h163 IrArray::Index EmitBlockIndex(llvm::Type* index_ty);
166 IrArray::Index GetTileIndexForBlockOrigin(const IrArray::Index& block_index);
169 IrArray::Index GetElementIndexForTileOrigin(const IrArray::Index& tile_index);
174 IrArray::Index GetUnnormalizedIndex(
175 const IrArray::Index& normalized_shape_index,
Dtuple_ops.cc36 void EmitTupleSelect(const IrArray& select, const IrArray& pred, in EmitTupleSelect()
72 void EmitTuple(const IrArray& tuple, absl::Span<llvm::Value* const> operands, in EmitTuple()
84 void EmitTuple(const IrArray& tuple, absl::Span<const IrArray> buffers, in EmitTuple()
90 [](const llvm_ir::IrArray& buffer) { return buffer.GetBasePointer(); }); in EmitTuple()
Dfused_ir_emitter.cc41 using llvm_ir::IrArray;
45 [=](const IrArray::Index& index) -> StatusOr<llvm::Value*> { in DefaultAction()
83 indexed_generators_[constant] = [=](const IrArray::Index& index) { in HandleConstant()
97 return IrArray(shape_constant, constant->shape()) in HandleConstant()
131 [=](const IrArray::Index& index) -> StatusOr<llvm::Value*> { in HandleGetTupleElement()
135 return IrArray(tuple_element_ptr, get_tuple_element->shape()) in HandleGetTupleElement()
146 [=](const IrArray::Index& index) -> llvm::Value* { in HandleParameter()
177 [=](const IrArray::Index& index) -> StatusOr<llvm::Value*> { in HandleTuple()
Dsort_util.cc145 const IrArray::Index& tiled_keys_index, int64 dimension_to_sort, in EmitTiledCompareLoop()
147 const std::vector<IrArray>& params, in EmitTiledCompareLoop()
194 IrArray::Index keys_index(keys_multi_index, params[i].GetShape(), in EmitTiledCompareLoop()
272 IrArray::Index keys_index(keys_multi_index, params[i].GetShape(), in EmitTiledCompareLoop()
293 int64 dimension_to_sort, const std::vector<IrArray>& values_arrays, in EmitSortInPlace()
340 [&](const IrArray::Index& tiles_index) -> Status { in EmitSortInPlace()
360 IrArray::Index keys_index(keys_multi_index, values_arrays[0].GetShape(), in EmitSortInPlace()
369 IrArray::Index keys_index(keys_multi_index, in EmitSortInPlace()
377 IrArray::Index keys_index(keys_multi_index, in EmitSortInPlace()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Ddot_op_emitter.cc110 const llvm_ir::IrArray& target_array,
111 const llvm_ir::IrArray& lhs_array,
112 const llvm_ir::IrArray& rhs_array,
113 const llvm_ir::IrArray* addend_array,
194 const llvm_ir::IrArray& target_array_;
195 const llvm_ir::IrArray& lhs_array_;
196 const llvm_ir::IrArray& rhs_array_;
197 const llvm_ir::IrArray* addend_array_;
206 const llvm_ir::IrArray& target_array, in DotOpEmitter()
207 const llvm_ir::IrArray& lhs_array, in DotOpEmitter()
[all …]
Ddot_op_emitter.h61 const llvm_ir::IrArray& target_array,
62 const llvm_ir::IrArray& lhs_array,
63 const llvm_ir::IrArray& rhs_array,
64 const llvm_ir::IrArray* addend_array,
Dir_emitter.h63 std::function<std::vector<llvm_ir::IrArray>()>;
125 const llvm_ir::IrArray::Index& index);
131 const llvm_ir::IrArray::Index& index);
137 const llvm_ir::IrArray::Index& index);
229 llvm_ir::IrArray GetIrArrayFor(const HloInstruction* hlo);
232 std::vector<llvm_ir::IrArray> GetIrArraysForOperandsOf(
242 llvm_ir::IrArray* array) { in AddAliasingInformationToIrArray()
382 const llvm_ir::IrArray& containing_array);
398 const llvm_ir::IrArray::Index& output_index,
414 const llvm_ir::IrArray& target_array,
[all …]
Delemental_ir_emitter.cc30 using xla::llvm_ir::IrArray;
113 const IrArray::Index& index) -> StatusOr<llvm::Value*> { in MakeElementGenerator()
124 return [this, hlo, &operand_to_generator](const IrArray::Index& index) { in MakeElementGenerator()
130 return [this, hlo, &operand_to_generator](const IrArray::Index& index) { in MakeElementGenerator()
137 return [this, hlo, &operand_to_generator](const IrArray::Index& index) { in MakeElementGenerator()
Dir_emitter.cc374 llvm_ir::EmitTuple(llvm_ir::IrArray(data_address, data_shape), in HandleInfeed()
643 const llvm_ir::IrArray::Index& index) { in EmitElementalReduceWindow()
662 const llvm_ir::IrArray::Index window_index = loops.AddLoopsForShape( in EmitElementalReduceWindow()
715 llvm_ir::IrArray::Index input_index(input_multi_index, operand->shape(), in EmitElementalReduceWindow()
785 [this, init_value](const llvm_ir::IrArray::Index& target_index) { in HandleSelectAndScatter()
792 const llvm_ir::IrArray::Index source_index = in HandleSelectAndScatter()
815 const llvm_ir::IrArray::Index window_index = window_loops.AddLoopsForShape( in HandleSelectAndScatter()
849 [&](const llvm_ir::IrArray::Index& operand_index) { in HandleSelectAndScatter()
856 llvm_ir::IrArray operand_array(GetIrArrayFor(operand)); in HandleSelectAndScatter()
857 llvm_ir::IrArray::Index operand_index( in HandleSelectAndScatter()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter.h72 std::function<std::vector<llvm_ir::IrArray>()>;
122 llvm_ir::IrArray GetIrArray(const HloInstruction& inst,
134 std::vector<llvm_ir::IrArray> ConstructIrArrayForOutputs(
190 std::vector<llvm_ir::IrArray> ir_arrays; in GetGeneratorForOperandIrArrays()
220 const llvm_ir::IrArray::Index& keys_index,
221 const llvm_ir::IrArray::Index& compare_keys_index,
222 const llvm_ir::IrArray& keys_array);
Dir_emitter_unnested.h57 std::function<void(const llvm_ir::IrArray::Index& output_tile_origin,
114 HloInstruction* hlo, const llvm_ir::IrArray::Index& index,
211 const HloInstruction* unnested_hlo, const llvm_ir::IrArray::Index& index,
257 const llvm_ir::IrArray::Index& index,
264 const llvm_ir::IrArray::Index& index,
271 const llvm_ir::IrArray::Index& index,
294 std::vector<llvm_ir::IrArray> ConstructIrArrayForInputs(
304 const std::vector<llvm_ir::IrArray>& param_arrays,
308 std::vector<llvm_ir::IrArray>* param_in_reduced_shape_arrays);
Dir_emitter_unnested.cc96 std::function<void(const llvm_ir::IrArray::Index& index, llvm::Value* y_loc,
105 using llvm_ir::IrArray;
662 IrArray output_array = GetIrArray(*fusion, *fusion); in HandleFusion()
707 const HloInstruction* unnested_hlo, const IrArray::Index& index, in EmitExtraOutputsForReduce()
834 auto loop_body_emitter = [=](const IrArray::Index& source_index) -> Status { in HandleSelectAndScatter()
858 const IrArray::Index window_index = window_loops.AddLoopsForShape( in HandleSelectAndScatter()
892 const auto save_operand_index = [&](const IrArray::Index& operand_index) { in HandleSelectAndScatter()
899 IrArray operand_array = GetIrArray(*operand, *select_and_scatter); in HandleSelectAndScatter()
900 IrArray::Index operand_index(operand_multi_index, operand->shape(), in HandleSelectAndScatter()
952 IrArray::Index selected_index(selected_multi_index, in HandleSelectAndScatter()
[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.cc42 absl::Span<const llvm_ir::IrArray> target_arrays, in ParallelLoopEmitter()
51 const llvm_ir::IrArray& target_array, in ParallelLoopEmitter()
58 std::vector<llvm_ir::IrArray::Index>
74 std::vector<llvm_ir::IrArray::Index> array_indices; in EmitIndexAndSetExitBasicBlock()
/external/tensorflow/tensorflow/compiler/xla/service/
Delemental_ir_emitter.h168 const llvm_ir::IrArray::Index& index);
173 const llvm_ir::IrArray::Index& index);
178 const llvm_ir::IrArray::Index& target_index);
183 const llvm_ir::IrArray::Index& index);
188 const llvm_ir::IrArray::Index& index);
193 const llvm_ir::IrArray::Index& index);
198 const llvm_ir::IrArray::Index& padded_index);
203 const llvm_ir::IrArray::Index& dot_result_index);
232 const llvm_ir::IrArray::Index& index, llvm::Value* raw_value);
Delemental_ir_emitter.cc52 using llvm_ir::IrArray;
1371 const llvm_ir::IrArray::Index& index, llvm::Value* raw_value) { in ConvertValueForDistribution()
1629 return [=](const llvm_ir::IrArray::Index& index) -> StatusOr<llvm::Value*> { in MakePhiloxRngElementGenerator()
1671 const llvm_ir::IrArray::Index& index) { in EmitElementalSelect()
1685 const llvm_ir::IrArray::Index& index) { in EmitElementalClamp()
1708 const llvm_ir::IrArray::Index& target_index) { in EmitElementalConcatenate()
1778 llvm_ir::IrArray::Index operand_index(operand_multi_index, operand->shape(), in EmitElementalConcatenate()
1815 const llvm_ir::IrArray::Index& index) { in EmitElementalDynamicSlice()
1826 llvm_ir::IrArray::Index zero_index(index_type); in EmitElementalDynamicSlice()
1854 llvm_ir::IrArray::Index input_index(input_multi_index, input_hlo->shape(), in EmitElementalDynamicSlice()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/cpu/tests/
Dcpu_noalias_test.cc84 llvm_ir::IrArray::Index zero2D({zero, zero}); in TEST_F()
92 llvm_ir::IrArray param_x_array(param_x_val, param_shape); in TEST_F()
102 llvm_ir::IrArray concat1_array(concat1_val, shape); in TEST_F()
112 llvm_ir::IrArray concat2_array(concat2_val, shape); in TEST_F()

12