Home
last modified time | relevance | path

Searched refs:ShapeIndex (Results 1 – 25 of 99) sorted by relevance

1234

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_input_output_alias_config.h47 Alias(AliasKind kind, int64 parameter_number, ShapeIndex parameter_index) in Alias()
54 ShapeIndex parameter_index;
66 Status SetUpAlias(const ShapeIndex& output_index, int64 param_number,
67 const ShapeIndex& param_index, AliasKind kind);
72 const ShapeIndex& param_index) const;
77 const ShapeIndex& param_index) const { in ParameterHasAlias()
82 bool OutputHasAlias(const ShapeIndex& output_index) const;
94 absl::optional<ShapeIndex> GetAliasedOutput(
95 int64 param_number, const ShapeIndex& param_index) const;
101 const ShapeIndex& output_index) const;
[all …]
Doptimize_input_output_buffer_alias_test.cc53 [&](const ShapeIndex&, const HloInputOutputAliasConfig::Alias&) { in AliasCount() argument
94 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); in TEST_F()
95 EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{1}); in TEST_F()
96 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{2}); in TEST_F()
97 EXPECT_EQ(config_.GetAliasedOutput(0, {3}), ShapeIndex{3}); in TEST_F()
110 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); in TEST_F()
111 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); in TEST_F()
124 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{3}); in TEST_F()
125 EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{2}); in TEST_F()
126 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); in TEST_F()
[all …]
Dtuple_points_to_analysis.h78 const ShapeIndex& index) const;
85 void AddPointedToBuffer(const LogicalBuffer& buffer, const ShapeIndex& index);
107 const SourceSet& tuple_sources(const ShapeIndex& index) const;
110 void add_tuple_source(const ShapeIndex& index, HloInstruction* tuple);
115 const BufferList& element(const ShapeIndex& index) const { in element()
118 BufferList* mutable_element(const ShapeIndex& index) { in mutable_element()
125 tree_.ForEachElement([&fn](const ShapeIndex& index, const Elem& elem) { in ForEachElement()
131 tree_.ForEachMutableElement([&fn](const ShapeIndex& index, Elem* elem) { in ForEachMutableElement()
138 [&fn](const ShapeIndex& index, const Elem& elem) { in ForEachElementWithStatus()
160 BufferAlias(HloInstruction* instruction, const ShapeIndex& index) in BufferAlias()
[all …]
Dhlo_input_output_alias_config.cc22 const ShapeIndex& output_index) const { in OutputHasAlias()
26 Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, in SetUpAlias()
28 const ShapeIndex& param_index, in SetUpAlias()
57 [&](const ShapeIndex& index, const absl::optional<Alias>& data) { in ToProto()
88 ShapeIndex output_index(entry.output_shape_index().begin(), in CreateFromProto()
91 ShapeIndex param_index(entry.parameter_shape_index().begin(), in CreateFromProto()
109 ForEachAlias([&](const ShapeIndex& output_index, const Alias& alias) { in ToString()
121 int64 param_number, const ShapeIndex& param_index) const { in ParameterAliasKind()
124 [&](const xla::ShapeIndex&, absl::optional<Alias> alias) { in ParameterAliasKind() argument
133 absl::optional<ShapeIndex> HloInputOutputAliasConfig::GetAliasedOutput( in GetAliasedOutput()
[all …]
Dtuple_points_to_analysis.cc51 [&ambiguous](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsAmbiguous() argument
60 ForEachElement([&](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsDistinct() argument
80 [&flat_set](const ShapeIndex& /*index*/, const BufferList& buffers) { in CreateFlattenedSet() argument
88 ForEachElement([&found, &buffer](const ShapeIndex& /*index*/, in ContainsBuffer()
98 const ShapeIndex& index) const { in ContainsBufferAtIndex() argument
104 const ShapeIndex& index) { in AddPointedToBuffer()
112 const ShapeIndex& index) const { in tuple_sources()
116 void PointsToSet::add_tuple_source(const ShapeIndex& index, in add_tuple_source()
187 const ShapeIndex& index,
204 [this, hlo_instruction](const ShapeIndex& index, in DefaultAction()
[all …]
Dhlo_dataflow_analysis.h86 const ShapeIndex& index = {}) const;
93 const ShapeIndex& index = {}) const;
95 const ShapeIndex& index = {});
106 const ShapeIndex& index = {}) const;
110 const ShapeIndex& index = {});
115 const ShapeIndex& index = {}) const {
119 const ShapeIndex& index = {}) {
144 const ShapeIndex& index,
152 const ShapeIndex& operand_index,
154 const ShapeIndex& user_index) const;
[all …]
Ddynamic_dimension_inference.cc80 HloInstruction* operand, ShapeIndex index, int64 dimension,
100 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in DefaultAction()
112 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleGetTupleElement()
115 ShapeIndex new_index = in HandleGetTupleElement()
125 hlo, [&](HloInstruction*, ShapeIndex index, int64 dimension, in HandleTuple()
135 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleBroadcast()
145 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandlePad()
168 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleReduce()
202 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleDot()
253 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleTranspose()
[all …]
Dhlo_liveness_analysis.cc51 using VisitorFunction = std::function<void(const ShapeIndex& /*index*/)>;
55 index_tree.ForEachElement([&](const ShapeIndex& shape_index, bool live) { in ForEachLiveIndex()
67 const ShapeIndex& shape_index, in MarkLiveAtIndex()
99 [&](const Shape& sub_shape, const ShapeIndex& shape_index) { in MarkLiveAtAllIndices()
126 ForEachLiveIndex(index_tree, [&](const ShapeIndex& shape_index) { in PropagateLivenessThroughTuple()
134 ShapeIndex operand_shape_index; in PropagateLivenessThroughTuple()
157 ForEachLiveIndex(index_tree, [&](const ShapeIndex& shape_index) { in PropagateLivenessThroughGTE()
158 ShapeIndex operand_shape_index(shape_index); in PropagateLivenessThroughGTE()
176 ForEachLiveIndex(index_tree, [&](const ShapeIndex& shape_index) { in PropagateLivenessThroughWhile()
207 ForEachLiveIndex(index_tree, [&](const ShapeIndex& shape_index) { in PropagateLivenessToParameterCallers()
[all …]
Dhlo_dataflow_analysis.cc79 const ShapeIndex& index) const { in ValueIsDefinedAt()
88 const HloInstruction* instruction, const ShapeIndex& index) const { in GetValueDefinedAt()
94 const HloInstruction* instruction, const ShapeIndex& index) { in GetValueDefinedAt()
100 const ShapeIndex& index, in NewHloValue()
156 const ShapeIndex& index, in ToString()
198 const ShapeIndex& index = pair.first; in Phi()
288 const HloInstruction* instruction, const ShapeIndex& index) const { in GetValueSet()
293 const ShapeIndex& index) { in GetValueSet()
323 const ShapeIndex& operand_index = pair.first; in UpdateSendValueSet()
326 ShapeIndex index = {0}; in UpdateSendValueSet()
[all …]
Doptimize_input_output_buffer_alias.cc50 absl::flat_hash_map<int64, std::queue<ShapeIndex>> size_to_input_index; in Build()
52 input_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
63 output_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
72 const ShapeIndex& input_index = it->second.front(); in Build()
73 const ShapeIndex& output_index = index; in Build()
Dhlo_input_output_alias_config_test.cc38 void expect_aliased(const ShapeIndex& output_index, int64 param_number, in expect_aliased()
39 const ShapeIndex& param_index, in expect_aliased()
41 absl::optional<ShapeIndex> aliased_output = in expect_aliased()
55 void expect_not_aliased(const ShapeIndex& output_index, int64 param_number, in expect_not_aliased()
56 const ShapeIndex& param_index, in expect_not_aliased()
58 absl::optional<ShapeIndex> aliased_output = in expect_not_aliased()
Dbfloat16_propagation.h167 const ShapeIndex& index) const;
172 const ShapeIndex& index) const;
183 const ShapeIndex& index,
200 absl::flat_hash_map<HloInstruction*, absl::flat_hash_map<Shape*, ShapeIndex>>
Dbfloat16_propagation.cc53 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineFusionComputationPrecision()
101 inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in RevertIfFusionInternalBF16Changes()
178 const Shape& subshape, const ShapeIndex& index) { in DetermineWhileComputationsPrecision()
207 const ShapeIndex& index) const { in AllUsersConsumeBF16()
281 ShapeIndex use_output_index{use.operand_number}; in AllUsersConsumeBF16()
290 ShapeIndex use_output_index; in AllUsersConsumeBF16()
342 const ShapeIndex& index) { in DetermineInstructionPrecision()
378 [hlo, this](const Shape& /* subshape */, const ShapeIndex& index) { in DetermineInstructionPrecision()
418 const ShapeIndex& index) { in AdjustCalledComputationParameters()
458 const ShapeIndex& index) { in AdjustCalledComputationRoot()
[all …]
Dbuffer_assignment.h131 const ShapeIndex& param_shape_index() const { in param_shape_index()
281 ShapeIndex param_shape_index, in set_entry_computation_parameter()
323 ShapeIndex param_shape_index_;
378 const HloInstruction* hlo, const ShapeIndex& shape_index) const;
383 const HloInstruction* instruction, const ShapeIndex& index) const;
388 const ShapeIndex& index) const;
398 const HloInstruction* instruction, const ShapeIndex& index) const;
411 const HloInstruction* instruction, const ShapeIndex& index) const { in GetSourceBuffers()
420 const ShapeIndex& shape_index_a,
422 const ShapeIndex& shape_index_b) const;
Dlogical_buffer_analysis.h39 const ShapeIndex& index) const;
56 void NewLogicalBuffer(HloInstruction* instruction, const ShapeIndex& index);
74 std::pair<const HloInstruction*, const ShapeIndex> p) const { in operator()
85 std::unordered_map<std::pair<const HloInstruction*, const ShapeIndex>,
Dlogical_buffer.h37 LogicalBuffer(HloInstruction* instruction, const ShapeIndex& index, Id id);
45 const ShapeIndex& index() const override { return index_; } in index()
58 ShapeIndex index_;
Dhlo_value.h41 ShapeIndex index;
72 ShapeIndex operand_index;
104 HloValue(Id id, HloInstruction* instruction, const ShapeIndex& index,
130 const ShapeIndex& defining_index() const { return defining_position().index; } in defining_index()
132 const ShapeIndex& index() const override { return defining_index(); } in index()
Ddynamic_parameter_binding_test.cc75 EXPECT_EQ(param->parameter_index, ShapeIndex({})); in TEST_F()
114 EXPECT_EQ(param->parameter_index, ShapeIndex({0})); in TEST_F()
157 EXPECT_EQ(param->parameter_index, ShapeIndex({0})); in TEST_F()
167 EXPECT_EQ(param2->parameter_index, ShapeIndex({0})); in TEST_F()
Dshaped_buffer_test.cc121 [&](const xla::ShapeIndex& index, se::DeviceMemoryBase* buffer) { in TEST()
130 xla::ShapeIndex subtree_index = {1}; in TEST()
133 output.buffers().ForEachElement([&](const xla::ShapeIndex& sub_index, in TEST()
135 xla::ShapeIndex orig_index = subtree_index; in TEST()
142 [&](const xla::ShapeIndex& index, const se::DeviceMemoryBase& buffer) { in TEST()
/external/tensorflow/tensorflow/compiler/xla/
Dliteral.h78 absl::Span<const NativeT> data(const ShapeIndex& shape_index = {}) const;
83 const ShapeIndex& shape_index = {}) const;
88 const void* untyped_data(const ShapeIndex& shape_index = {}) const;
89 int64 size_bytes(const ShapeIndex& shape_index = {}) const;
114 const ShapeIndex& shape_index) const;
128 const ShapeIndex& shape_index = {}) const;
132 const ShapeIndex& shape_index = {}) const;
142 int64 sparse_element_number, const ShapeIndex& shape_index = {}) const;
150 const ShapeIndex& shape_index = {}) const;
216 int64 element_count(const ShapeIndex& index = {}) const {
[all …]
Dshape_tree.h45 std::pair<ShapeIndex, T> data;
49 explicit ShapeTreeNode(ShapeIndex index) in ShapeTreeNode()
51 ShapeTreeNode(ShapeIndex index, T data) in ShapeTreeNode()
152 std::pair<ShapeIndex, T>>;
156 const std::pair<ShapeIndex, T>>;
268 ForEachElement([&](const ShapeIndex& index, const T& t) { in Map()
277 ForEachMutableElement([&](const ShapeIndex& index, T* t) { in Map()
290 const ShapeIndex& source_base_index,
291 const ShapeIndex& target_base_index);
416 ShapeIndex shape_index = node->data.first; in InitChildren()
[all …]
Dshape_util.h65 class ShapeIndex {
67 ShapeIndex() = default;
68 ShapeIndex(std::initializer_list<int64> init) : indices_(init) {} in ShapeIndex() function
70 ShapeIndex(InputIt start, InputIt end) : indices_(start, end) {} in ShapeIndex() function
95 bool operator==(const ShapeIndex& other) const {
98 bool operator!=(const ShapeIndex& other) const { return !(*this == other); }
99 bool operator<(const ShapeIndex& other) const {
106 friend H AbslHashValue(H h, const ShapeIndex& index) { in AbslHashValue()
123 ShapeIndexView(const ShapeIndex& shape_index, int64 offset = 0)
150 ShapeIndex ToShapeIndex() const { return ShapeIndex(begin(), end()); } in ToShapeIndex()
[all …]
Dshape_tree_test.cc65 int_tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TestShapeConstructor() argument
74 [&num_nodes](const ShapeIndex& /*index*/, bool data) { in TestShapeConstructor() argument
91 tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TestInitValueConstructor() argument
99 [&num_nodes](const ShapeIndex& /*index*/, int* data) { in TestInitValueConstructor() argument
107 tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TestInitValueConstructor() argument
164 [&sum](const ShapeIndex& /*index*/, int data) { sum += data; }); in TEST_F() argument
176 [](const ShapeIndex& /*index*/, int* data) { *data = 0; }); in TEST_F() argument
383 t.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { in TEST_F() argument
452 std::vector<ShapeIndex> v; in TEST_F()
456 EXPECT_EQ(v, (std::vector<ShapeIndex>{{}, in TEST_F()
[all …]
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.cc256 const xla::ShapeIndex& buffer_index) { in DiscardAllocation()
296 XRTTupleAllocation* parent, const xla::ShapeIndex& subshape, in MakeSubBuffer()
314 [](const xla::ShapeIndex& index, in MakeSubBuffer()
321 [&](const xla::ShapeIndex& index, XRTBufferAllocationPtr* buffer) { in MakeSubBuffer()
324 xla::ShapeIndex parent_index = subshape; in MakeSubBuffer()
351 [&](const xla::ShapeIndex& index, const ExpandedTupleInput& element) { in ExpandTreeOfTuples()
409 [&](const xla::ShapeIndex& index, const ExpandedTupleInput& element) { in MakeTuple()
442 elements.ForEachElement([&](const xla::ShapeIndex& index, in MakeTuple()
454 [&](const xla::ShapeIndex& index, XRTBufferAllocationPtr* buffer) { in MakeTuple()
462 [](const xla::ShapeIndex& index, in MakeTuple()
[all …]
Dxrt_state.h106 const xla::ShapeIndex& subshape,
173 void DiscardAllocation(const xla::ShapeIndex& buffer_index);
182 const xla::ShapeIndex& source_index,
183 const xla::ShapeIndex& dest_index);
190 const std::function<bool(const xla::ShapeIndex&)>& release_checker);

1234