Home
last modified time | relevance | path

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

12345

/external/tensorflow/tensorflow/compiler/xla/service/
Doptimize_input_output_buffer_alias_test.cc48 [&](const ShapeIndex&, const HloInputOutputAliasConfig::Alias&) { in AliasCount() argument
96 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); in TEST_F()
97 EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{1}); in TEST_F()
98 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{2}); in TEST_F()
99 EXPECT_EQ(config_.GetAliasedOutput(0, {3}), ShapeIndex{3}); in TEST_F()
112 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); in TEST_F()
113 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); in TEST_F()
126 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{3}); in TEST_F()
127 EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{2}); in TEST_F()
128 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); in TEST_F()
[all …]
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 …]
Dtuple_points_to_analysis.cc52 [&ambiguous](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsAmbiguous() argument
61 ForEachElement([&](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsDistinct() argument
81 [&flat_set](const ShapeIndex& /*index*/, const BufferList& buffers) { in CreateFlattenedSet() argument
89 ForEachElement([&found, &buffer](const ShapeIndex& /*index*/, in ContainsBuffer()
99 const ShapeIndex& index) const { in ContainsBufferAtIndex() argument
105 const ShapeIndex& index) { in AddPointedToBuffer()
113 const ShapeIndex& index) const { in tuple_sources()
117 void PointsToSet::add_tuple_source(const ShapeIndex& index, in add_tuple_source()
188 const ShapeIndex& index,
205 [this, hlo_instruction](const ShapeIndex& index, in DefaultAction()
[all …]
Dtuple_points_to_analysis.h79 const ShapeIndex& index) const;
86 void AddPointedToBuffer(const LogicalBuffer& buffer, const ShapeIndex& index);
108 const SourceSet& tuple_sources(const ShapeIndex& index) const;
111 void add_tuple_source(const ShapeIndex& index, HloInstruction* tuple);
116 const BufferList& element(const ShapeIndex& index) const { in element()
119 BufferList* mutable_element(const ShapeIndex& index) { in mutable_element()
126 tree_.ForEachElement([&fn](const ShapeIndex& index, const Elem& elem) { in ForEachElement()
132 tree_.ForEachMutableElement([&fn](const ShapeIndex& index, Elem* elem) { in ForEachMutableElement()
139 [&fn](const ShapeIndex& index, const Elem& elem) { in ForEachElementWithStatus()
161 BufferAlias(HloInstruction* instruction, const ShapeIndex& index) in BufferAlias()
[all …]
Dhlo_dataflow_analysis.h53 const ShapeIndex& user_index)>;
82 const ShapeIndex& index = {}) const;
89 const ShapeIndex& index = {}) const;
91 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,
[all …]
Dhlo_input_output_alias_config.cc23 const ShapeIndex& output_index) const { in OutputHasAlias()
27 Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, in SetUpAlias()
29 const ShapeIndex& param_index, in SetUpAlias()
58 [&](const ShapeIndex& index, const absl::optional<Alias>& data) { in ToProto()
89 ShapeIndex output_index(entry.output_shape_index().begin(), in CreateFromProto()
92 ShapeIndex param_index(entry.parameter_shape_index().begin(), in CreateFromProto()
114 ForEachAlias([&](const ShapeIndex& output_index, const Alias& alias) { in ToString()
126 int64 param_number, const ShapeIndex& param_index) const { in ParameterAliasKind()
129 [&](const xla::ShapeIndex&, absl::optional<Alias> alias) { in ParameterAliasKind() argument
138 absl::optional<ShapeIndex> HloInputOutputAliasConfig::GetAliasedOutput( in GetAliasedOutput()
[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 …]
Ddynamic_dimension_inference.cc106 HloInstruction* operand, ShapeIndex index, int64 dimension,
130 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in DefaultAction()
143 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleGetTupleElement()
147 ShapeIndex new_index = in HandleGetTupleElement()
158 hlo, [&](HloInstruction*, ShapeIndex index, int64 dimension, in HandleTuple()
170 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleBroadcast()
190 ShapeIndex data_output = {0}; in HandleCustomCall()
198 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleCustomCall()
214 hlo, [&](HloInstruction* operand, ShapeIndex index, in HandleSort()
232 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandlePad()
[all …]
Doptimize_input_output_buffer_alias.cc55 ShapeIndex index; in Build()
61 *input_shapes[i], [&](const Shape& subshape, const ShapeIndex& index) { in Build()
72 output_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
79 const ShapeIndex& input_index = entry.index; in Build()
80 const ShapeIndex& output_index = index; in Build()
Dhlo_input_output_alias_config_test.cc37 void expect_aliased(const ShapeIndex& output_index, int64 param_number, in expect_aliased()
38 const ShapeIndex& param_index, in expect_aliased()
40 absl::optional<ShapeIndex> aliased_output = in expect_aliased()
54 void expect_not_aliased(const ShapeIndex& output_index, int64 param_number, in expect_not_aliased()
55 const ShapeIndex& param_index, in expect_not_aliased()
57 absl::optional<ShapeIndex> aliased_output = in expect_not_aliased()
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()
210 const ShapeIndex& index = pair.first; in Phi()
317 const HloInstruction* instruction, const ShapeIndex& index) const { in GetValueSet()
322 const ShapeIndex& index) { in GetValueSet()
366 const ShapeIndex& operand_index = pair.first; in UpdateSendValueSet()
369 ShapeIndex index = {0}; in UpdateSendValueSet()
[all …]
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>>
Dlogical_buffer.h37 LogicalBuffer(HloInstruction* instruction, const ShapeIndex& index, Id id);
45 const ShapeIndex& index() const override { return index_; } in index()
58 ShapeIndex index_;
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()
364 const ShapeIndex& index) { in DetermineInstructionPrecision()
394 [hlo, this](const Shape& /* subshape */, const ShapeIndex& index) { in DetermineInstructionPrecision()
434 const ShapeIndex& index) { in AdjustCalledComputationParameters()
474 const ShapeIndex& index) { in AdjustCalledComputationRoot()
[all …]
Dlogical_buffer_analysis.h40 const ShapeIndex& index) const;
57 void NewLogicalBuffer(HloInstruction* instruction, const ShapeIndex& index);
77 std::pair<const HloInstruction*, const ShapeIndex> p) const { in operator()
88 std::unordered_map<std::pair<const HloInstruction*, const ShapeIndex>,
Dhlo_value.h42 ShapeIndex index;
73 ShapeIndex operand_index;
105 HloValue(Id id, HloInstruction* instruction, const ShapeIndex& index,
131 const ShapeIndex& defining_index() const { return defining_position().index; } in defining_index()
133 const ShapeIndex& index() const override { return defining_index(); } in index()
Ddynamic_parameter_binding_test.cc74 EXPECT_EQ(param->parameter_index, ShapeIndex({})); in TEST_F()
113 EXPECT_EQ(param->parameter_index, ShapeIndex({0})); in TEST_F()
156 EXPECT_EQ(param->parameter_index, ShapeIndex({0})); in TEST_F()
166 EXPECT_EQ(param2->parameter_index, ShapeIndex({0})); in TEST_F()
Dshaped_buffer_test.cc124 [&](const xla::ShapeIndex& index, se::DeviceMemoryBase* buffer) { in TEST()
133 xla::ShapeIndex subtree_index = {1}; in TEST()
136 output.buffers().ForEachElement([&](const xla::ShapeIndex& sub_index, in TEST()
138 xla::ShapeIndex orig_index = subtree_index; in TEST()
145 [&](const xla::ShapeIndex& index, const se::DeviceMemoryBase& buffer) { in TEST()
162 [&](const xla::ShapeIndex& index, se::DeviceMemoryBase* buffer) { in TEST()
Dbuffer_assignment.h134 const ShapeIndex& param_shape_index() const { in param_shape_index()
284 ShapeIndex param_shape_index, in set_entry_computation_parameter()
326 ShapeIndex param_shape_index_;
385 const HloInstruction* hlo, const ShapeIndex& shape_index) const;
390 const HloInstruction* instruction, const ShapeIndex& index) const;
395 const ShapeIndex& index) const;
405 const HloInstruction* instruction, const ShapeIndex& index) const;
418 const HloInstruction* instruction, const ShapeIndex& index) const { in GetSourceBuffers()
427 const ShapeIndex& shape_index_a,
429 const ShapeIndex& shape_index_b) const;
/external/tensorflow/tensorflow/compiler/xla/
Dshape_tree.h45 std::pair<ShapeIndex, T> data;
49 explicit ShapeTreeNode(ShapeIndex index) in ShapeTreeNode()
51 ShapeTreeNode(ShapeIndex index, T data) in ShapeTreeNode()
153 std::pair<ShapeIndex, T>>;
157 const std::pair<ShapeIndex, T>>;
269 ForEachElement([&](const ShapeIndex& index, const T& t) { in Map()
278 ForEachMutableElement([&](const ShapeIndex& index, T* t) { in Map()
291 const ShapeIndex& source_base_index,
292 const ShapeIndex& target_base_index);
294 StatusOr<ShapeTree<T>> SubShapeTree(const ShapeIndex& index) const;
[all …]
Dshape_util.h67 class ShapeIndex {
69 ShapeIndex() = default;
70 ShapeIndex(std::initializer_list<int64> init) : indices_(init) {} in ShapeIndex() function
72 ShapeIndex(InputIt start, InputIt end) : indices_(start, end) {} in ShapeIndex() function
74 explicit ShapeIndex(ShapeIndexView v);
99 bool operator==(const ShapeIndex& other) const {
102 bool operator!=(const ShapeIndex& other) const { return !(*this == other); }
103 bool operator<(const ShapeIndex& other) const {
110 friend H AbslHashValue(H h, const ShapeIndex& index) { in AbslHashValue()
127 ShapeIndexView(const ShapeIndex& shape_index, int64 offset = 0)
[all …]
Dliteral.h77 absl::Span<const NativeT> data(const ShapeIndex& shape_index = {}) const;
82 const void* untyped_data(const ShapeIndex& shape_index = {}) const;
83 int64 size_bytes(const ShapeIndex& shape_index = {}) const;
108 const ShapeIndex& shape_index) const;
122 const ShapeIndex& shape_index = {}) const;
229 int64 element_count(const ShapeIndex& index = {}) const {
274 const ShapeIndex& shape_index = {}) const;
384 ShapeIndex index; in ForEachSubpiece()
386 [&func](const ShapeIndex& index, const Piece& piece) { in ForEachSubpiece()
398 ShapeIndex index; in ForEachSubpieceWithStatus()
[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.h116 const xla::ShapeIndex& subshape,
205 void DiscardAllocation(const xla::ShapeIndex& buffer_index);
214 const xla::ShapeIndex& source_index,
215 const xla::ShapeIndex& dest_index);
232 const std::function<xla::StatusOr<bool>(const xla::ShapeIndex&)>&
Dxrt_state.cc344 const xla::ShapeIndex& buffer_index) { in DiscardAllocation()
363 XRTTupleAllocation* parent, const xla::ShapeIndex& subshape, in MakeSubBuffer()
381 [](const xla::ShapeIndex& index, in MakeSubBuffer()
388 [&](const xla::ShapeIndex& index, XRTBufferAllocationPtr* buffer) { in MakeSubBuffer()
391 xla::ShapeIndex parent_index = subshape; in MakeSubBuffer()
425 [&](const xla::ShapeIndex& index, const ExpandedTupleInput& element) { in ExpandTreeOfTuples()
483 [&](const xla::ShapeIndex& index, const ExpandedTupleInput& element) { in MakeTuple()
516 elements.ForEachElement([&](const xla::ShapeIndex& index, in MakeTuple()
528 [&](const xla::ShapeIndex&, XRTBufferAllocationPtr* buffer) { in MakeTuple() argument
534 [](const xla::ShapeIndex& index, in MakeTuple()
[all …]

12345