Skip to content

Commit 086e0ff

Browse files
author
Gautham Ganapathy
committed
Removed ExtendedTensor constructors that required passing in a graph instance
Summary: REF T67791 Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, samuelh Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, samuelh Maniphest Tasks: T67791 Differential Revision: https://phabricator.sourcevertex.net/D73973
1 parent 15f0d4a commit 086e0ff

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

45 files changed

+190
-299
lines changed

‎tensorflow/compiler/plugin/poplar/driver/extended_graph.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ class ExtendedGraph : public poplar::Graph {
102102
const poplar::Type& type, poplar::ArrayRef<std::size_t> shape,
103103
const T* values, const poplar::DebugContext& debugContext = {"<const>"}) {
104104
auto tensor = poplar::Graph::addConstant(type, shape, values, debugContext);
105-
return {std::move(tensor), *this};
105+
return {std::move(tensor)};
106106
}
107107

108108
template <typename T>
@@ -111,7 +111,7 @@ class ExtendedGraph : public poplar::Graph {
111111
const poplar::DebugContext& debugContext = {
112112
"<const>"}) {
113113
auto tensor = poplar::Graph::addConstant(type, shape, val, debugContext);
114-
return {tensor, *this};
114+
return tensor;
115115
}
116116

117117
ExtendedTensor addConstantHalf(

‎tensorflow/compiler/plugin/poplar/driver/extended_tensor.h‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,6 @@ class ExtendedTensor : public poplar::Tensor {
3838
ExtendedTensor(const poplar::Tensor& tensor) // NOLINT
3939
: poplar::Tensor(tensor) {}
4040

41-
ExtendedTensor(const poplar::Tensor& tensor,
42-
const ExtendedGraph& graph) // NOLINT
43-
: poplar::Tensor(tensor) {}
44-
45-
ExtendedTensor(poplar::Graph& graph) // NOLINT
46-
: poplar::Tensor() {}
4741

4842
ExtendedTensor reshape(poplar::ArrayRef<std::size_t> shape) const;
4943

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/execution_counter.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,7 @@ class ExecutionCounterOp : public PoplarOpDef {
4343
auto counter_copy = poputil::duplicate(
4444
graph, counter, seq, {debug_info, "ExecutionCounterCopy"});
4545

46-
TF_RETURN_IF_ERROR(AddOutputTensor(tensor_map, inst, 0,
47-
DriverTensor(counter_copy, graph)));
46+
TF_RETURN_IF_ERROR(AddOutputTensor(tensor_map, inst, 0, counter_copy));
4847

4948
return seq;
5049
}

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/poplin/cholesky.cc‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ class CholeskyOp : public PoplarOpDef {
6363
auto out = poplin::cholesky(graph, a, lower, seq, {dnai, "Cholesky"},
6464
poplar_options, &res.planning_cache);
6565

66-
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, DriverTensor(out, graph)));
66+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, DriverTensor(out)));
6767

6868
return seq;
6969
}
@@ -87,11 +87,9 @@ class CholeskyOp : public PoplarOpDef {
8787
TF_ASSIGN_OR_RETURN(poplar::OptionFlags poplar_options,
8888
GetCholeskyOptionsForInst(inst, res));
8989

90-
auto out = DriverTensor(
91-
poplin::createCholeskyInput(graph, type, shape, lower,
92-
{debug_context, "createInput"},
93-
poplar_options, &res.planning_cache),
94-
graph);
90+
DriverTensor out = poplin::createCholeskyInput(
91+
graph, type, shape, lower, {debug_context, "createInput"},
92+
poplar_options, &res.planning_cache);
9593

9694
MappingHelper::RemapTensor(res.linear_mapping_state, graph, out);
9795

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/poplin/triangular_solve.cc‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ class TriangularSolveOp : public PoplarOpDef {
9898
{debug_name_and_id, "TriangularSolve"},
9999
poplar_options, &res.planning_cache);
100100

101-
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, DriverTensor(x, graph)));
101+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, DriverTensor(x)));
102102
return seq;
103103
}
104104

@@ -145,19 +145,15 @@ class TriangularSolveOp : public PoplarOpDef {
145145
DriverTensor out;
146146
switch (input_index) {
147147
case 0: {
148-
out = DriverTensor(poplin::createTriangularSolveInputLHS(
149-
graph, type_a, type_b, poplar_shape_a,
150-
poplar_shape_b, left_side, {debug_info, "lhs"},
151-
poplar_options, &res.planning_cache),
152-
graph);
148+
out = poplin::createTriangularSolveInputLHS(
149+
graph, type_a, type_b, poplar_shape_a, poplar_shape_b, left_side,
150+
{debug_info, "lhs"}, poplar_options, &res.planning_cache);
153151
break;
154152
}
155153
case 1: {
156-
out = DriverTensor(poplin::createTriangularSolveInputRHS(
157-
graph, type_a, type_b, poplar_shape_a,
158-
poplar_shape_b, left_side, {debug_info, "rhs"},
159-
poplar_options, &res.planning_cache),
160-
graph);
154+
out = poplin::createTriangularSolveInputRHS(
155+
graph, type_a, type_b, poplar_shape_a, poplar_shape_b, left_side,
156+
{debug_info, "rhs"}, poplar_options, &res.planning_cache);
161157
break;
162158
}
163159
default:

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/poplin/weights_transpose_chans_flip_xy.cc‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,10 @@ StatusOr<DriverTensor> AddConvWeightsTransposeChansFlipXY(
4444
TF_ASSIGN_OR_RETURN(poplar::OptionFlags opts,
4545
GetConvolutionOptionsForInst(inst, resources));
4646

47-
DriverTensor out_weights = DriverTensor(
48-
poplin::createWeights(
49-
graph, conv_params,
50-
{debug_name_and_id, "createWeights_TransposeChansFlipXY"}, opts,
51-
&resources.planning_cache),
52-
graph);
47+
DriverTensor out_weights = poplin::createWeights(
48+
graph, conv_params,
49+
{debug_name_and_id, "createWeights_TransposeChansFlipXY"}, opts,
50+
&resources.planning_cache);
5351

5452
out_weights = RemoveGroupsDimensionFromWeights(conv_params, out_weights);
5553

@@ -91,10 +89,9 @@ class WeightsTransposeChansFlipXYOp : public PoplarOpDef {
9189
in_weights = AddGroupsDimensionToWeights(conv_params, in_weights,
9290
/* swap_features= */ true);
9391

94-
DriverTensor out_weights = DriverTensor(
92+
DriverTensor out_weights =
9593
poplin::createWeights(graph, conv_params, {debug_info, "CreateWeights"},
96-
opts, &res.planning_cache),
97-
graph);
94+
opts, &res.planning_cache);
9895

9996
poplin::weightsTransposeChansFlipXY(
10097
graph, in_weights, out_weights, seq,
@@ -106,7 +103,7 @@ class WeightsTransposeChansFlipXYOp : public PoplarOpDef {
106103
conv_dimension_numbers, out_weights, /* swap_features = */ true);
107104

108105
TF_CHECK_OK(
109-
AddOutputTensor(tensor_map, inst, 0, DriverTensor(out_weights, graph)));
106+
AddOutputTensor(tensor_map, inst, 0, DriverTensor(out_weights)));
110107

111108
return seq;
112109
}

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/popnn/arg_min_max.cc‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ class ArgMinMaxOp : public PoplarOpDef {
5555
indices = indices.reshape(output_dimensions);
5656

5757
TF_RETURN_IF_ERROR(
58-
AddOutputTensor(tensor_map, inst, 0, DriverTensor(indices, graph)));
58+
AddOutputTensor(tensor_map, inst, 0, DriverTensor(indices)));
5959
return Status::OK();
6060
}
6161

@@ -129,9 +129,9 @@ class MaxMinAndArgMinMaxOp : public ArgMinMaxOp {
129129
values = values.reshape(output_dimensions);
130130

131131
TF_RETURN_IF_ERROR(
132-
AddOutputTensor(tensor_map, inst, 0, DriverTensor(values, graph)));
132+
AddOutputTensor(tensor_map, inst, 0, DriverTensor(values)));
133133
TF_RETURN_IF_ERROR(
134-
AddOutputTensor(tensor_map, inst, 1, DriverTensor(indices, graph)));
134+
AddOutputTensor(tensor_map, inst, 1, DriverTensor(indices)));
135135
return Status::OK();
136136
}
137137
};

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/popnn/ctc_loss.cc‎

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -88,10 +88,8 @@ class CTCLossOpBase : public PoplarOpDef {
8888
label_lengths.reinterpret(poplar::UNSIGNED_INT), seq, blank_index,
8989
*plan, debug_info);
9090

91-
TF_CHECK_OK(
92-
AddOutputTensor(tensor_map, inst, 0, DriverTensor(loss, graph)));
93-
TF_CHECK_OK(
94-
AddOutputTensor(tensor_map, inst, 1, DriverTensor(grad, graph)));
91+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0, DriverTensor(loss)));
92+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 1, DriverTensor(grad)));
9593

9694
return seq;
9795
}
@@ -129,17 +127,15 @@ class CTCLossOpBase : public PoplarOpDef {
129127
const int64_t max_time = ShapeUtil::GetDimension(data_shape, 0);
130128
const int64_t num_classes = ShapeUtil::GetDimension(data_shape, 2);
131129
return DriverTensor(popnn::ctc::createDataInput(
132-
graph, dtype, batch_size, max_time, num_classes,
133-
*plan, {debug_info, "data"}),
134-
graph);
130+
graph, dtype, batch_size, max_time, num_classes, *plan,
131+
{debug_info, "data"}));
135132
}
136133
case 1: {
137134
const int64_t max_label_length =
138135
ShapeUtil::GetDimension(labels_shape, 1);
139136
return DriverTensor(popnn::ctc::createLabelsInput(
140-
graph, dtype, batch_size, max_label_length,
141-
*plan, {debug_info, "labels"}),
142-
graph);
137+
graph, dtype, batch_size, max_label_length, *plan,
138+
{debug_info, "labels"}));
143139
}
144140
default: {
145141
return FailedPrecondition(
@@ -236,12 +232,12 @@ class CTCBeamSearchOpBase : public PoplarOpDef {
236232
PerformBeamSearch(graph, data, data_lengths, seq, blank_index,
237233
beam_width, top_paths, *plan, debug_info);
238234

239-
TF_CHECK_OK(AddOutputTensor(
240-
tensor_map, inst, 0, DriverTensor(outputs.label_probabilities, graph)));
235+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 0,
236+
DriverTensor(outputs.label_probabilities)));
241237
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 1,
242-
DriverTensor(outputs.label_lengths, graph)));
238+
DriverTensor(outputs.label_lengths)));
243239
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, 2,
244-
DriverTensor(outputs.decoded_labels, graph)));
240+
DriverTensor(outputs.decoded_labels)));
245241
return seq;
246242
}
247243

@@ -275,9 +271,8 @@ class CTCBeamSearchOpBase : public PoplarOpDef {
275271
dtype.toString(), in_dtype.toString());
276272
}
277273
return DriverTensor(popnn::ctc_infer::createDataInput(
278-
graph, dtype, batch_size, max_time, num_classes,
279-
*plan, {debug_info, "data"}),
280-
graph);
274+
graph, dtype, batch_size, max_time, num_classes, *plan,
275+
{debug_info, "data"}));
281276
}
282277
};
283278

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/popnn/gru.cc‎

Lines changed: 13 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -106,9 +106,8 @@ class GRULayerBaseOp : public PoplarOpDef {
106106
const HloInstruction* inst,
107107
TensorMap& tensor_map, bool training) {
108108
for (int64_t j = 0; j < OutputTensorCount(); ++j) {
109-
TF_CHECK_OK(
110-
AddOutputTensor(tensor_map, inst, j,
111-
DriverTensor(args[j + InputTensorCount()], graph)));
109+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, j,
110+
DriverTensor(args[j + InputTensorCount()])));
112111
}
113112
}
114113

@@ -127,17 +126,13 @@ class GRULayerBaseOp : public PoplarOpDef {
127126
switch (input_index) {
128127
case 0: {
129128
// Allocate GRU input tensor.
130-
return DriverTensor(
131-
popnn::gru::createInput(graph, gru_params, {debug_info}, gru_opts,
132-
&res.planning_cache),
133-
graph);
129+
return DriverTensor(popnn::gru::createInput(
130+
graph, gru_params, {debug_info}, gru_opts, &res.planning_cache));
134131
}
135132
case 1: {
136133
// Allocate initial state tensor.
137-
return DriverTensor(
138-
popnn::gru::createInitialState(graph, gru_params, {debug_info},
139-
gru_opts, &res.planning_cache),
140-
graph);
134+
return DriverTensor(popnn::gru::createInitialState(
135+
graph, gru_params, {debug_info}, gru_opts, &res.planning_cache));
141136
}
142137
case 2: {
143138
// Allocate GRU weights kernel.
@@ -146,28 +141,22 @@ class GRULayerBaseOp : public PoplarOpDef {
146141
std::tie(input_weights, output_weights) =
147142
popnn::gru::createWeightsKernel(graph, gru_params, {debug_info},
148143
gru_opts, &res.planning_cache);
149-
return DriverTensor(PackGruKernel(input_weights, output_weights),
150-
graph);
144+
return DriverTensor(PackGruKernel(input_weights, output_weights));
151145
}
152146
case 3: {
153147
// Allocate GRU weights biases.
154-
return DriverTensor(
155-
popnn::gru::createWeightsBiases(graph, gru_params, {debug_info},
156-
gru_opts, &res.planning_cache),
157-
graph);
148+
return DriverTensor(popnn::gru::createWeightsBiases(
149+
graph, gru_params, {debug_info}, gru_opts, &res.planning_cache));
158150
}
159151
case 4: {
160152
// Allocate AUGRU seq_len
161153
return DriverTensor(popops::createSliceableTensor(
162-
graph, poplar::INT, {gru_params.rnn.batchSize},
163-
{0}, {1}, 0, name),
164-
graph);
154+
graph, poplar::INT, {gru_params.rnn.batchSize}, {0}, {1}, 0, name));
165155
}
166156
case 5: {
167157
// Allocate AUGRU attention
168158
return DriverTensor(popnn::gru::createAttention(graph, gru_params, name)
169-
.dimShuffle({1, 0}),
170-
graph);
159+
.dimShuffle({1, 0}));
171160
}
172161
default: {
173162
return xla::FailedPrecondition(
@@ -275,9 +264,8 @@ class GRULayerFwdOp : public GRULayerBaseOp {
275264
int total_param_count = InputTensorCount() + OutputTensorCount();
276265
GRULayerBaseOp::SetOutputTensor(graph, args, inst, tensor_map, training);
277266
if (training) {
278-
TF_CHECK_OK(
279-
AddOutputTensor(tensor_map, inst, OutputTensorCount(),
280-
DriverTensor(args[total_param_count], graph)));
267+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, OutputTensorCount(),
268+
DriverTensor(args[total_param_count])));
281269
}
282270
}
283271

‎tensorflow/compiler/plugin/poplar/driver/ops/custom_ops/popnn/lstm.cc‎

Lines changed: 14 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -100,9 +100,8 @@ class LstmLayerBaseOp : public PoplarOpDef {
100100
const HloInstruction* inst,
101101
TensorMap& tensor_map, bool training) {
102102
for (int64_t j = 0; j < OutputTensorCount(); ++j) {
103-
TF_CHECK_OK(
104-
AddOutputTensor(tensor_map, inst, j,
105-
DriverTensor(args[j + InputTensorCount()], graph)));
103+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, j,
104+
DriverTensor(args[j + InputTensorCount()])));
106105
}
107106
}
108107

@@ -121,24 +120,18 @@ class LstmLayerBaseOp : public PoplarOpDef {
121120
switch (input_index) {
122121
case 0: {
123122
// Allocate LSTM input tensor
124-
return DriverTensor(
125-
popnn::lstm::createInput(graph, lstm_params, {debug_info},
126-
lstm_opts, &res.planning_cache),
127-
graph);
123+
return DriverTensor(popnn::lstm::createInput(
124+
graph, lstm_params, {debug_info}, lstm_opts, &res.planning_cache));
128125
}
129126
case 1: {
130127
// Allocate initial output (h) tensor
131-
return DriverTensor(
132-
popnn::lstm::createInitialOutput(graph, lstm_params, {debug_info},
133-
lstm_opts, &res.planning_cache),
134-
graph);
128+
return DriverTensor(popnn::lstm::createInitialOutput(
129+
graph, lstm_params, {debug_info}, lstm_opts, &res.planning_cache));
135130
}
136131
case 2: {
137132
// Allocate initial cell state (c) tensor
138133
return DriverTensor(popnn::lstm::createInitialCellState(
139-
graph, lstm_params, {debug_info}, lstm_opts,
140-
&res.planning_cache),
141-
graph);
134+
graph, lstm_params, {debug_info}, lstm_opts, &res.planning_cache));
142135
}
143136
case 3: {
144137
// Allocate LSTM weights kernel
@@ -147,22 +140,18 @@ class LstmLayerBaseOp : public PoplarOpDef {
147140
std::tie(input_weights, output_weights) =
148141
popnn::lstm::createWeightsKernel(graph, lstm_params, {debug_info},
149142
lstm_opts, &res.planning_cache);
150-
return DriverTensor(PackLstmKernel(input_weights, output_weights),
151-
graph);
143+
return DriverTensor(PackLstmKernel(input_weights, output_weights));
152144
}
153145
case 4: {
154146
// Allocate LSTM weights biases
155-
return DriverTensor(
156-
popnn::lstm::createWeightsBiases(graph, lstm_params, {debug_info},
157-
lstm_opts, &res.planning_cache),
158-
graph);
147+
return DriverTensor(popnn::lstm::createWeightsBiases(
148+
graph, lstm_params, {debug_info}, lstm_opts, &res.planning_cache));
159149
}
160150
case 5: {
161151
// Allocate DynamicLSTM seq_len
162152
return DriverTensor(popops::createSliceableTensor(
163-
graph, poplar::INT, {lstm_params.rnn.batchSize},
164-
{0}, {1}, 0, name),
165-
graph);
153+
graph, poplar::INT, {lstm_params.rnn.batchSize}, {0}, {1}, 0,
154+
name));
166155
}
167156
default: {
168157
return xla::FailedPrecondition(
@@ -274,9 +263,8 @@ class LstmLayerFwdOp : public LstmLayerBaseOp {
274263
const int64_t total_param_count = InputTensorCount() + OutputTensorCount();
275264
LstmLayerBaseOp::SetOutputTensor(graph, args, inst, tensor_map, training);
276265
if (training) {
277-
TF_CHECK_OK(
278-
AddOutputTensor(tensor_map, inst, OutputTensorCount(),
279-
DriverTensor(args[total_param_count], graph)));
266+
TF_CHECK_OK(AddOutputTensor(tensor_map, inst, OutputTensorCount(),
267+
DriverTensor(args[total_param_count])));
280268
}
281269
}
282270

0 commit comments

Comments
 (0)