experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Remove use of designated initializers in dynamic dimension inference.

Browse Source 
PiperOrigin-RevId: 286330416
Change-Id: I5878ab23b9596ac0c5dbf91394ee6ba8e5683c77
This commit is contained in:
Yunxing Dai 2019-12-18 22:18:58 -08:00 committed by Saleem Abdulrasool
parent 4dfee0db43
commit 4c0721928b
2 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/compiler/xla/service/dynamic_dimension_inference.cc
View File

@ -469,7 +469,7 @@ Status DynamicDimensionInferenceVisitor::HandleSetDimensionSize(
// Propagate dynamic dimension indicated by this set dimension size // Propagate dynamic dimension indicated by this set dimension size
// instruction. // instruction.
parent_->SetDynamicSize(hlo, {}, hlo->dimension(), hlo->mutable_operand(1), parent_->SetDynamicSize(hlo, {}, hlo->dimension(), hlo->mutable_operand(1),
{.stride = 1, .multiple_of = 1}); DimensionConstraint(1, 1));
} }
// Also Propagate dynamic dimension already set by operands. // Also Propagate dynamic dimension already set by operands.
@ -813,7 +813,7 @@ Status DynamicDimensionInferenceVisitor::HandleReshape(HloInstruction* hlo) {
parent_->SetDynamicSize( parent_->SetDynamicSize(
reshape, {}, output_dynamic_dimension, new_dynamic_size, reshape, {}, output_dynamic_dimension, new_dynamic_size,
{.stride = 1, .multiple_of = constraint.multiple_of / divisor}); DimensionConstraint(1, constraint.multiple_of / divisor));
} }
if (input_dim_size < output_dim_size) { if (input_dim_size < output_dim_size) {
@ -850,12 +850,12 @@ Status DynamicDimensionInferenceVisitor::HandleReshape(HloInstruction* hlo) {
hlo->parent()->AddInstruction(HloInstruction::CreateBinary( hlo->parent()->AddInstruction(HloInstruction::CreateBinary(
output_dynamic_size->shape(), HloOpcode::kMultiply, output_dynamic_size->shape(), HloOpcode::kMultiply,
new_dynamic_size, operand_dynamic_size)); new_dynamic_size, operand_dynamic_size));
int64 new_multiple_of_constraint =
constraint.multiple_of * output_dim_size /
operand->shape().dimensions(input_dynamic_dimension);
parent_->SetDynamicSize( parent_->SetDynamicSize(
reshape, {}, output_dynamic_dimension, new_dynamic_size, reshape, {}, output_dynamic_dimension, new_dynamic_size,
{.stride = 1, DimensionConstraint(1, new_multiple_of_constraint));
.multiple_of =
constraint.multiple_of * output_dim_size /
operand->shape().dimensions(input_dynamic_dimension)});
} }
return Status::OK(); return Status::OK();
@ -1227,7 +1227,7 @@ Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) {
parent_->SetDynamicSize(target_parameter, parent_->SetDynamicSize(target_parameter,
dynamic_dimension.parameter_index, dynamic_dimension.parameter_index,
dynamic_dimension.dimension, dynamic_size, dynamic_dimension.dimension, dynamic_size,
{.stride = 1, .multiple_of = 1}); DimensionConstraint(1, 1));
return Status::OK(); return Status::OK();
}); });
} }

3
tensorflow/compiler/xla/service/dynamic_dimension_inference.h
View File

@ -149,6 +149,9 @@ class DynamicDimensionInference {
// //
// //
struct DimensionConstraint { struct DimensionConstraint {
explicit DimensionConstraint(int64 s, int64 m)
: stride(s), multiple_of(m) {}
DimensionConstraint() : stride(1), multiple_of(1) {}
// Stride represents the distance of a newly placed element and the previous // Stride represents the distance of a newly placed element and the previous
// placed element on this dynamic dimension. // placed element on this dynamic dimension.
int64 stride; int64 stride;