Home
last modified time | relevance | path

Searched refs:shape_index (Results 1 – 25 of 45) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_liveness_analysis.cc55 index_tree.ForEachElement([&](const ShapeIndex& shape_index, bool live) { in ForEachLiveIndex() argument
57 func(shape_index); in ForEachLiveIndex()
67 const ShapeIndex& shape_index, in MarkLiveAtIndex() argument
77 if (it->second.element(shape_index) == false) { in MarkLiveAtIndex()
79 *it->second.mutable_element(shape_index) = true; in MarkLiveAtIndex()
81 << " shape_index: " << shape_index.ToString(); in MarkLiveAtIndex()
99 [&](const Shape& sub_shape, const ShapeIndex& shape_index) { in MarkLiveAtAllIndices() argument
100 if (it->second.element(shape_index) == false) { in MarkLiveAtAllIndices()
102 *it->second.mutable_element(shape_index) = true; in MarkLiveAtAllIndices()
104 << " shape_index: " << shape_index.ToString(); in MarkLiveAtAllIndices()
[all …]
Dhlo_cost_analysis.cc836 [this, fusion](const Shape& subshape, const ShapeIndex& shape_index) { in HandleFusion() argument
840 if (shape_index.empty()) { in HandleFusion()
846 SetOutputBytesAccessed(shape_index, size); in HandleFusion()
849 } else if (shape_index.size() == 1) { in HandleFusion()
852 ->operand(shape_index[0]) in HandleFusion()
855 ->operand(shape_index[0]) in HandleFusion()
859 SetOutputBytesAccessed(shape_index, size); in HandleFusion()
864 SetOutputBytesAccessed(shape_index, GetShapeSize(subshape)); in HandleFusion()
875 const ShapeIndex& shape_index) { in HandleFusion() argument
877 current_properties_.find(GetOutputBytesAccessedKey(shape_index)); in HandleFusion()
[all …]
Dhlo_liveness_analysis.h52 const ShapeIndex& shape_index) const;
Dlayout_assignment.cc769 [&](const Shape& subshape, const ShapeIndex& shape_index) { in CheckParameterLayout() argument
770 if (!ShapeUtil::IsLeafIndex(parameter_layout.shape(), shape_index) || in CheckParameterLayout()
776 ShapeUtil::GetSubshape(parameter->shape(), shape_index))) { in CheckParameterLayout()
1327 const auto shape_index = in PropagateOperandConstraint() local
1333 user, shape_index)); in PropagateOperandConstraint()
1410 [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateOperandConstraint() argument
1424 .InstructionDefinesBufferAtIndex(user, shape_index)) { in PropagateOperandConstraint()
1432 user, shape_index)); in PropagateOperandConstraint()
1451 user->shape(), [&](const Shape& subshape, const ShapeIndex& shape_index) { in PropagateOperandConstraint() argument
1459 user, shape_index)) { in PropagateOperandConstraint()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dhlo_to_ir_bindings.h52 ShapeIndexView shape_index = {});
72 ShapeIndexView shape_index = {}) const {
75 return it->second.element(shape_index);
87 const ShapeIndex& shape_index = {});
98 ShapeIndexView shape_index,
Dhlo_to_ir_bindings.cc141 ShapeIndexView shape_index, in GetTypedIrValue() argument
144 ShapeUtil::GetSubshape(hlo.shape(), shape_index), ir_value, b_); in GetTypedIrValue()
156 ShapeIndexView shape_index) { in BindHloToIrValue() argument
160 llvm::Value* typed_ir_value = GetTypedIrValue(hlo, shape_index, ir_value); in BindHloToIrValue()
166 *(base_ptrs_[&hlo].mutable_element(shape_index)) = typed_ir_value; in BindHloToIrValue()
193 const ShapeIndex& shape_index) { in GetIrArray() argument
194 llvm::Value* base_ptr = GetBasePointer(hlo, shape_index); in GetIrArray()
196 << "Buffer not assigned for shape_index " << shape_index.ToString() in GetIrArray()
199 ShapeUtil::GetSubshape(hlo.shape(), shape_index)); in GetIrArray()
Dir_emitter.h127 const ShapeIndex& shape_index = {}) {
128 return bindings_.GetIrArray(inst, consumer, shape_index);
132 ShapeIndexView shape_index = {}) const {
133 return bindings_.GetBasePointer(inst, shape_index);
Dinfeed_thunk.cc47 const ShapeIndex& shape_index = source.first; in ExecuteOnStream() local
50 ShapeUtil::GetSubshape(source_buffers.shape(), shape_index); in ExecuteOnStream()
Doutfeed_thunk.cc60 const ShapeIndex& shape_index = output_leaf_it->first; in ExecuteOnStream() local
79 ShapeUtil::GetSubshape(output_buffers->shape(), shape_index); in ExecuteOnStream()
Dgpu_compiler.cc1057 ShapeIndex shape_index; in GetMlirAllocationInfo() local
1060 shape_index.push_back(i.getSExtValue()); in GetMlirAllocationInfo()
1063 if (*rank != shape_index.size()) { in GetMlirAllocationInfo()
1067 rank.emplace(shape_index.size()); in GetMlirAllocationInfo()
1069 auto& o = (*output_info)[shape_index]; in GetMlirAllocationInfo()
1076 if (shape_index.size() > 1) { in GetMlirAllocationInfo()
1081 if (shape_index.empty()) { in GetMlirAllocationInfo()
1084 if (shape_index[0] >= output_shapes.size()) { in GetMlirAllocationInfo()
1085 output_shapes.resize(shape_index[0] + 1); in GetMlirAllocationInfo()
1087 output_shapes[shape_index[0]] = TypeToShape(arg.getType()); in GetMlirAllocationInfo()
/external/tensorflow/tensorflow/compiler/xla/
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;
290 const ShapeIndex& shape_index = {}) const;
559 const Piece& piece(const ShapeIndex& shape_index) const { in piece() argument
561 for (const auto i : shape_index) { in piece()
593 absl::Span<NativeT> data(const ShapeIndex& shape_index = {});
[all …]
Dliteral.cc234 const ShapeIndex& shape_index) const { in GetDynamicSize()
235 return piece(shape_index).GetDynamicSize(dim_index); in GetDynamicSize()
533 const ShapeIndex& shape_index, in SetDynamicSize() argument
535 Shape* subshape_ = ShapeUtil::GetMutableSubshape(shape_.get(), shape_index); in SetDynamicSize()
542 piece(shape_index).SetDynamicSize(dim_index, size); in SetDynamicSize()
709 const ShapeIndex& shape_index) const { in Relayout()
712 Shape* subshape = ShapeUtil::GetMutableSubshape(&new_shape, shape_index); in Relayout()
973 const ShapeIndex& shape_index) const { in GetAsString()
974 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetAsString()
978 return Get<bool>(multi_index, shape_index) ? "true" : "false"; in GetAsString()
[all …]
Dliteral_comparison.cc283 const ShapeIndex& shape_index, ErrorSpec error, in Compare() argument
286 NearComparator<NativeT> comparator(expected, actual, shape_index, error, in Compare()
319 const ShapeIndex& shape_index, ErrorSpec error, in NearComparator() argument
324 shape_index_(shape_index), in NearComparator()
678 const ShapeIndex& shape_index, in EqualHelper() argument
684 ShapeIndex next_index = shape_index; in EqualHelper()
763 shape_index); in EqualHelper()
774 const ShapeIndex& shape_index, const ErrorSpec& error, in NearHelper() argument
784 ShapeIndex element_index = shape_index; in NearHelper()
801 if (!return_status.ok() && shape_index.empty()) { in NearHelper()
[all …]
Dshape_tree.h458 ShapeIndex shape_index = node->data.first; in InitChildren() local
459 shape_index.push_back(0); in InitChildren()
469 shape_index[shape_index.size() - 1] = i; in InitChildren()
475 nodes_.emplace_back(shape_index, init_value); in InitChildren()
494 ShapeIndex shape_index = node->data.first; in InitChildren() local
495 shape_index.push_back(0); in InitChildren()
505 shape_index[shape_index.size() - 1] = i; in InitChildren()
511 nodes_.emplace_back(shape_index); in InitChildren()
Dshape_layout.cc76 ShapeIndexView shape_index) { in ResetLayout() argument
78 *ShapeUtil::GetMutableSubshape(&shape_, shape_index)->mutable_layout() = in ResetLayout()
Dshape_util.h128 ShapeIndexView(const ShapeIndex& shape_index, int64 offset = 0)
129 : indices_(shape_index.data() + offset, shape_index.size() - offset) {
130 CHECK_LE(offset, shape_index.size());
176 std::ostream& operator<<(std::ostream& out, const ShapeIndex& shape_index);
177 std::ostream& operator<<(std::ostream& out, const ShapeIndexView& shape_index);
Dliteral_util.cc64 [&](const Shape& subshape, const ShapeIndex& shape_index) { in ConvertType() argument
68 auto src = literal.data<FromNativeT>(shape_index); in ConvertType()
69 auto dest = result.data<ToNativeT>(shape_index); in ConvertType()
75 /*dest_shape_index=*/shape_index, in ConvertType()
76 /*src_shape_index=*/shape_index)); in ConvertType()
Dliteral_comparison.h40 const LiteralSlice& mismatches, const ShapeIndex& shape_index)>;
Dshape_layout.h82 void ResetLayout(const Layout& layout, ShapeIndexView shape_index);
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/
Dtpu_dynamic_padding_mapper.mlir8 // shape_index: 2
30 // shape_index: 2
35 // shape_index: 3
57 // shape_index: 2
62 // shape_index: 3
67 // shape_index: 1
88 // shape_index: 2
110 // shape_index: 2
160 // shape_index: 2
182 // shape_index: 2
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dmhlo_to_lhlo_with_xla.cc448 xla::ShapeIndex* shape_index, OpBuilder* b, Location loc) { in RewriteFusionOperand() argument
452 shape_index->push_back(i); in RewriteFusionOperand()
454 auto v, RewriteFusionOperand(root, shape.tuple_shapes(i), shape_index, in RewriteFusionOperand()
457 shape_index->pop_back(); in RewriteFusionOperand()
462 GetOrCreateArrayView(root, shape, *shape_index)); in RewriteFusionOperand()
489 xla::ShapeIndex shape_index; in EmitFusionOp() local
491 auto arg, RewriteFusionOperand(operand, operand->shape(), &shape_index, in EmitFusionOp()
1331 const xla::ShapeIndex& shape_index) { in GetOrCreateArrayView() argument
1334 auto& cached_value = slices_[std::make_pair(instr, shape_index)]; in GetOrCreateArrayView()
1339 if (instr->IsConstant() && shape_index.empty()) { in GetOrCreateArrayView()
[all …]
/external/tensorflow/tensorflow/core/protobuf/tpu/
Ddynamic_padding.proto14 int32 shape_index = 2; field
/external/tensorflow/tensorflow/compiler/xla/tests/
Dhlo_test_base.h283 ShapeIndexView shape_index) { in ForceResultLayout() argument
286 ->ResetLayout(layout, shape_index); in ForceResultLayout()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc1644 int arg_index, int32_t shape_index, in AddDynamicParameterBindingEntry() argument
1648 entry->set_target_param_dim_num(shape_index); in AddDynamicParameterBindingEntry()
1713 const int32_t shape_index = shape_index_attr.getInt(); in AddDynamicParameterBindings() local
1714 if (arg_type && (shape_index < 0 || shape_index >= arg_type.getRank())) in AddDynamicParameterBindings()
1717 arg_type.getRank(), shape_index)); in AddDynamicParameterBindings()
1718 else if (shape_index < 0) in AddDynamicParameterBindings()
1720 PaddingMapNegativeShapeIndexMsg(i, element_index, shape_index)); in AddDynamicParameterBindings()
1722 if (!used_shape_indices.insert(shape_index).second) in AddDynamicParameterBindings()
1724 PaddingMapUniqueShapeIndexMsg(i, element_index, shape_index)); in AddDynamicParameterBindings()
1744 AddDynamicParameterBindingEntry(dynamic_parameter_binding, i, shape_index, in AddDynamicParameterBindings()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtpu_dynamic_padding_mapper.cc164 it.first.push_back(padding.shape_index()); in AnnotateFunctionArgumentsWithPaddings()

12