Home
last modified time | relevance | path

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

123456

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_input_output_alias_config.h45 Alias(int64 parameter_number, ShapeIndex parameter_index,
52 ShapeIndex parameter_index;
73 Status SetUpAlias(const ShapeIndex& output_index, int64 param_number,
74 const ShapeIndex& param_index,
80 const ShapeIndex& param_index) const { in ParameterHasAlias()
85 bool OutputHasAlias(const ShapeIndex& output_index) const;
97 absl::optional<ShapeIndex> GetAliasedOutput(
98 int64 param_number, const ShapeIndex& param_index) const;
104 const ShapeIndex& output_index) const;
109 const ShapeIndex& param_index) const;
[all …]
Doptimize_input_output_buffer_alias_test.cc48 [&](const ShapeIndex&, const HloInputOutputAliasConfig::Alias&) { in AliasCount() argument
89 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); in TEST_F()
90 EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{1}); in TEST_F()
91 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{2}); in TEST_F()
92 EXPECT_EQ(config_.GetAliasedOutput(0, {3}), ShapeIndex{3}); in TEST_F()
105 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{0}); in TEST_F()
106 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); in TEST_F()
119 EXPECT_EQ(config_.GetAliasedOutput(0, {0}), ShapeIndex{3}); in TEST_F()
120 EXPECT_EQ(config_.GetAliasedOutput(0, {1}), ShapeIndex{2}); in TEST_F()
121 EXPECT_EQ(config_.GetAliasedOutput(0, {2}), ShapeIndex{1}); in TEST_F()
[all …]
Dtuple_points_to_analysis.cc54 [&ambiguous](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsAmbiguous() argument
63 ForEachElement([&](const ShapeIndex& /*index*/, const BufferList& points_to) { in IsDistinct() argument
83 [&flat_set](const ShapeIndex& /*index*/, const BufferList& buffers) { in CreateFlattenedSet() argument
91 ForEachElement([&found, &buffer](const ShapeIndex& /*index*/, in ContainsBuffer()
101 const ShapeIndex& index) const { in ContainsBufferAtIndex() argument
107 const ShapeIndex& index) { in AddPointedToBuffer()
115 const ShapeIndex& index) const { in tuple_sources()
119 void PointsToSet::add_tuple_source(const ShapeIndex& index, in add_tuple_source()
190 const ShapeIndex& index,
207 [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.h60 const ShapeIndex& user_index)>;
90 const ShapeIndex& index = {}) const;
97 const ShapeIndex& index = {}) const;
99 const ShapeIndex& index = {});
114 const ShapeIndex& index = {}) const;
118 const ShapeIndex& index = {});
123 const ShapeIndex& index = {}) const {
127 const ShapeIndex& index = {}) {
152 const ShapeIndex& index,
160 const ShapeIndex& operand_index,
[all …]
Dhlo_input_output_alias_config.cc24 const ShapeIndex& output_index) const { in OutputHasAlias()
29 const ShapeIndex& output_index, int64 param_number, in SetUpAlias()
30 const ShapeIndex& param_index, in SetUpAlias()
57 [&](const ShapeIndex& index, const absl::optional<Alias>& data) { in ToProto()
83 ShapeIndex output_index(entry.output_shape_index().begin(), in CreateFromProto()
86 ShapeIndex param_index(entry.parameter_shape_index().begin(), in CreateFromProto()
103 ForEachAlias([&](const ShapeIndex& output_index, const Alias& alias) { in ToString()
115 const ShapeIndex& index = p.first; in ToShortString()
125 int64 param_number, const ShapeIndex& param_index) const { in ParameterMustAlias()
128 [&](const xla::ShapeIndex&, absl::optional<Alias> alias) { in ParameterMustAlias() argument
[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 …]
Dexecutable.h102 void SetBuffer(const ShapeIndex& index, MaybeOwningDeviceMemory buffer) { in SetBuffer()
106 void SetUnownedBuffer(const ShapeIndex& index,
109 void SetUnownedIndex(const ShapeIndex& index) { in SetUnownedIndex()
113 void ClearUnownedIndex(const ShapeIndex& index) { in ClearUnownedIndex()
117 const std::set<ShapeIndex>& unowned_indices() { return unowned_indices_; } in unowned_indices()
123 MaybeOwningDeviceMemory* MutableBuffer(const ShapeIndex& index) { in MutableBuffer()
127 const MaybeOwningDeviceMemory& Buffer(const ShapeIndex& index) const { in Buffer()
141 std::set<ShapeIndex> unowned_indices_;
175 void AddAliasedIndex(ShapeIndex index) { in AddAliasedIndex()
207 std::vector<ShapeIndex> ConsumeAliasedIndices() { in ConsumeAliasedIndices()
[all …]
Doptimize_input_output_buffer_alias.cc52 ShapeIndex index; in Build()
57 input_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
67 output_shape, [&](const Shape& subshape, const ShapeIndex& index) { in Build()
74 const ShapeIndex& input_index = entry.index; in Build()
75 const ShapeIndex& output_index = index; in Build()
Ddynamic_dimension_inference.h51 HloInstruction* GetDynamicSize(HloInstruction* inst, const ShapeIndex& index,
57 const ShapeIndex& index) const;
65 const ShapeIndex& index);
70 void SetDynamicSize(HloInstruction* inst, const ShapeIndex& index, int64 dim,
91 ShapeIndex index;
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.cc124 const ShapeIndex& index) const { in ValueIsDefinedAt()
133 const HloInstruction* instruction, const ShapeIndex& index) const { in GetValueDefinedAt()
139 const HloInstruction* instruction, const ShapeIndex& index) { in GetValueDefinedAt()
145 const ShapeIndex& index, in NewHloValue()
202 const ShapeIndex& index, in ToString()
256 const ShapeIndex& index = pair.first; in Phi()
369 const HloInstruction* instruction, const ShapeIndex& index) const { in GetValueSet()
374 const ShapeIndex& index) { in GetValueSet()
418 const ShapeIndex& operand_index = pair.first; in UpdateSendValueSet()
421 ShapeIndex index = {0}; in UpdateSendValueSet()
[all …]
Ddynamic_dimension_inference.cc155 HloInstruction* operand, ShapeIndex index, int64 dimension,
159 ShapeIndex index, int64 dimension, HloInstruction* dynamic_size)>;
202 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in DefaultAction()
214 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleGetTupleElement()
217 ShapeIndex new_index = in HandleGetTupleElement()
227 hlo, [&](HloInstruction*, ShapeIndex index, int64 dimension, in HandleTuple()
237 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleBroadcast()
255 ShapeIndex data_output = {0}; in HandleCustomCall()
290 hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, in HandleCustomCall()
345 [&](HloInstruction* operand, ShapeIndex index, int64 dynamic_dimension, in HandleSort()
[all …]
Dmemory_space_assignment_utils.cc68 if (position.index == ShapeIndex({0})) { in IsValueAllowedInAlternateMemory()
72 } else if (position.index == ShapeIndex({1})) { in IsValueAllowedInAlternateMemory()
83 if (position.index == ShapeIndex({0})) { in IsValueAllowedInAlternateMemory()
87 } else if (position.index == ShapeIndex({1})) { in IsValueAllowedInAlternateMemory()
Dbfloat16_propagation.h172 const ShapeIndex& index) const;
177 const ShapeIndex& index) const;
188 const ShapeIndex& index,
205 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_;
Dlogical_buffer_analysis.h40 const ShapeIndex& index) const;
57 void NewLogicalBuffer(HloInstruction* instruction, const ShapeIndex& index);
78 std::pair<const HloInstruction*, const ShapeIndex> p) const { in operator()
89 std::unordered_map<std::pair<const HloInstruction*, const ShapeIndex>,
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()
161 [&](const xla::ShapeIndex& index, se::DeviceMemoryBase* buffer) { in TEST()
Dbfloat16_propagation.cc55 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineFusionComputationPrecision()
103 inst->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in RevertIfFusionInternalBF16Changes()
180 const Shape& subshape, const ShapeIndex& index) { in DetermineWhileComputationsPrecision()
215 root->shape(), [&](const Shape& subshape, const ShapeIndex& index) { in DetermineConditionalComputationsPrecision()
236 const ShapeIndex& index) const { in AllUsersConsumeBF16()
318 ShapeIndex use_output_index{use.operand_number}; in AllUsersConsumeBF16()
327 ShapeIndex use_output_index; in AllUsersConsumeBF16()
411 const ShapeIndex& index) { in DetermineInstructionPrecision()
441 [hlo, this](const Shape& /* subshape */, const ShapeIndex& index) { in DetermineInstructionPrecision()
481 const ShapeIndex& index) { in AdjustCalledComputationParameters()
[all …]
Dlogical_buffer_analysis.cc87 const ShapeIndex& index) const { in GetBuffer()
92 const ShapeIndex& index) { in NewLogicalBuffer()
106 [this, hlo_instruction](const Shape& shape, const ShapeIndex& index) { in DefaultAction()
194 absl::flat_hash_set<ShapeIndex> aliased_outputs; in HandleCustomCall()
199 [&](const Shape& shape, const ShapeIndex& index) { in HandleCustomCall()
/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()
155 std::pair<ShapeIndex, T>>;
159 const std::pair<ShapeIndex, T>>;
166 std::pair<ShapeIndex, T>>;
170 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,
[all …]
Dshape_util.h68 class ShapeIndex {
70 ShapeIndex() = default;
71 ShapeIndex(std::initializer_list<int64> init) : indices_(init) {} in ShapeIndex() function
73 ShapeIndex(InputIt start, InputIt end) : indices_(start, end) {} in ShapeIndex() function
75 explicit ShapeIndex(ShapeIndexView v);
100 bool operator==(const ShapeIndex& other) const {
103 bool operator!=(const ShapeIndex& other) const { return !(*this == other); }
104 bool operator<(const ShapeIndex& other) const {
111 friend H AbslHashValue(H h, const ShapeIndex& index) { in AbslHashValue()
128 ShapeIndexView(const ShapeIndex& shape_index, int64 offset = 0)
[all …]
Dliteral.h78 absl::Span<const NativeT> data(const ShapeIndex& shape_index = {}) const;
83 const void* untyped_data(const ShapeIndex& shape_index = {}) const;
84 int64 size_bytes(const ShapeIndex& shape_index = {}) const;
117 const ShapeIndex& shape_index) const;
124 int32 GetDynamicSize(int64 dim_index, const ShapeIndex& shape_index) const;
138 const ShapeIndex& shape_index = {}) const;
245 int64 element_count(const ShapeIndex& index = {}) const {
290 const ShapeIndex& shape_index = {}) const;
424 ShapeIndex index; in ForEachSubpiece()
426 [&func](const ShapeIndex& index, const Piece& piece) { in ForEachSubpiece()
[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.h117 const xla::ShapeIndex& subshape,
206 void DiscardAllocation(const xla::ShapeIndex& buffer_index);
215 const xla::ShapeIndex& source_index,
216 const xla::ShapeIndex& dest_index);
232 const std::function<xla::StatusOr<bool>(const xla::ShapeIndex&)>&

123456