Home | History | Annotate | Download | only in xla 
      1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #include "tensorflow/compiler/xla/layout_util.h"
     17 
     18 #include <stddef.h>
     19 #include <algorithm>
     20 #include <functional>
     21 #include <random>
     22 #include <string>
     23 #include <unordered_map>
     24 #include <vector>
     25 
     26 #include "tensorflow/compiler/xla/protobuf_util.h"
     27 #include "tensorflow/compiler/xla/shape_util.h"
     28 #include "tensorflow/compiler/xla/status_macros.h"
     29 #include "tensorflow/compiler/xla/types.h"
     30 #include "tensorflow/compiler/xla/util.h"
     31 #include "tensorflow/core/lib/core/errors.h"
     32 #include "tensorflow/core/lib/strings/numbers.h"
     33 #include "tensorflow/core/lib/strings/str_util.h"
     34 #include "tensorflow/core/lib/strings/strcat.h"
     35 #include "tensorflow/core/platform/logging.h"
     36 #include "tensorflow/core/platform/protobuf.h"
     37 
     38 namespace xla {
     39 namespace {
     40 
     41 // Internal helper for GetDefaultLayoutForShape and SetToDefaultLayout. Sets
     42 // minor_to_major to the value that represents the default layout.
     43 void SetDefaultLayoutToContainer(
     44     tensorflow::protobuf::RepeatedField<tensorflow::protobuf_int64>*
     45         minor_to_major) {
     46   // The default XLA layout is major-to-minor (dim 0 is major).
     47   // For more information on XLA layouts, see:
     48   // https://www.tensorflow.org/performance/xla/shapes
     49   const int64 size = minor_to_major->size();
     50   for (int64 i = 0; i < size; ++i) {
     51     minor_to_major->Set(i, size - 1 - i);
     52   }
     53 }
     54 
     55 }  // namespace
     56 
     57 /* static */ Layout LayoutUtil::MakeLayout(
     58     tensorflow::gtl::ArraySlice<int64> minor_to_major) {
     59   Layout layout;
     60   layout.set_format(DENSE);
     61   for (int64 dimension_number : minor_to_major) {
     62     layout.add_minor_to_major(dimension_number);
     63   }
     64   return layout;
     65 }
     66 
     67 /* static */ Layout LayoutUtil::MakeSparseLayout(int64 max_sparse_elements) {
     68   Layout layout;
     69   layout.set_format(SPARSE);
     70   layout.set_max_sparse_elements(max_sparse_elements);
     71   return layout;
     72 }
     73 
     74 namespace {
     75 
     76 // Internal helper that creates a default layout for an array of the given rank.
     77 Layout CreateDefaultLayoutForRank(int64 rank) {
     78   Layout layout;
     79   layout.set_format(DENSE);
     80   tensorflow::protobuf::RepeatedField<tensorflow::protobuf_int64>*
     81       minor_to_major = layout.mutable_minor_to_major();
     82   minor_to_major->Resize(rank, 0);
     83   SetDefaultLayoutToContainer(minor_to_major);
     84   return layout;
     85 }
     86 
     87 }  // namespace
     88 
     89 /* static */ Layout LayoutUtil::GetDefaultLayoutForShape(const Shape& shape) {
     90   // A Layout proto corresponds to a single array, not a tuple.
     91   DCHECK(!ShapeUtil::IsTuple(shape));
     92   return CreateDefaultLayoutForRank(shape.dimensions_size());
     93 }
     94 
     95 /* static */ Layout LayoutUtil::GetDefaultLayoutForRank(int64 rank) {
     96   return CreateDefaultLayoutForRank(rank);
     97 }
     98 
     99 /* static */ Layout LayoutUtil::GetDefaultLayoutForR2() {
    100   return CreateDefaultLayoutForRank(2);
    101 }
    102 
    103 /* static */ Layout LayoutUtil::GetDefaultLayoutForR3() {
    104   return CreateDefaultLayoutForRank(3);
    105 }
    106 
    107 /* static */ Layout LayoutUtil::GetDefaultLayoutForR4() {
    108   return CreateDefaultLayoutForRank(4);
    109 }
    110 
    111 /* static */ void LayoutUtil::SetToDefaultLayout(Shape* shape) {
    112   if (ShapeUtil::IsTuple(*shape)) {
    113     // Tuple shape.
    114     for (auto& element_shape : *shape->mutable_tuple_shapes()) {
    115       SetToDefaultLayout(&element_shape);
    116     }
    117     shape->clear_layout();
    118   } else if (ShapeUtil::IsOpaque(*shape)) {
    119     shape->clear_layout();
    120   } else {
    121     shape->mutable_layout()->set_format(DENSE);
    122     tensorflow::protobuf::RepeatedField<tensorflow::protobuf_int64>*
    123         minor_to_major = shape->mutable_layout()->mutable_minor_to_major();
    124     minor_to_major->Resize(shape->dimensions_size(), 0);
    125     SetDefaultLayoutToContainer(minor_to_major);
    126   }
    127 }
    128 
    129 /* static */ Shape LayoutUtil::GetWithDefaultLayout(const Shape& shape) {
    130   Shape copy(shape);
    131   LayoutUtil::SetToDefaultLayout(&copy);
    132   return copy;
    133 }
    134 
    135 /* static */ void LayoutUtil::SetToDefaultLayout(ProgramShape* program_shape) {
    136   for (auto& parameter_shape : *program_shape->mutable_parameters()) {
    137     LayoutUtil::SetToDefaultLayout(&parameter_shape);
    138   }
    139   LayoutUtil::SetToDefaultLayout(program_shape->mutable_result());
    140 }
    141 
    142 /* static */ tensorflow::Status LayoutUtil::ValidateLayoutInShape(
    143     const Shape& shape) {
    144   if (ShapeUtil::IsTuple(shape)) {
    145     // Tuple shape.
    146     if (shape.has_layout()) {
    147       return InvalidArgument("tuple should not have a layout field");
    148     }
    149     for (auto& element_shape : shape.tuple_shapes()) {
    150       TF_RETURN_IF_ERROR(ValidateLayoutInShape(element_shape));
    151     }
    152     return tensorflow::Status::OK();
    153   } else if (ShapeUtil::IsOpaque(shape)) {
    154     if (shape.has_layout()) {
    155       return InvalidArgument("opaque should not have a layout field");
    156     }
    157     return tensorflow::Status::OK();
    158   } else {
    159     // Array shape.
    160     if (!shape.has_layout()) {
    161       return InvalidArgument("shape %s does not have a layout",
    162                              ShapeUtil::HumanString(shape).c_str());
    163     }
    164     return ValidateLayoutForShape(shape.layout(), shape);
    165   }
    166 }
    167 
    168 /* static */ tensorflow::Status LayoutUtil::ValidateLayoutForShape(
    169     const Layout& layout, const Shape& shape) {
    170   if (ShapeUtil::IsTuple(shape)) {
    171     return InvalidArgument("a single Layout is not valid for tuple shapes");
    172   }
    173 
    174   if (ShapeUtil::IsOpaque(shape)) {
    175     return tensorflow::Status::OK();
    176   }
    177 
    178   if (layout.format() == INVALID_FORMAT) {
    179     return InvalidArgument(
    180         "Layout does not have a valid format: layout {%s}, shape {%s}",
    181         layout.ShortDebugString().c_str(), shape.ShortDebugString().c_str());
    182   }
    183 
    184   if (layout.format() == DENSE) {
    185     if (layout.minor_to_major_size() != ShapeUtil::Rank(shape)) {
    186       return InvalidArgument(
    187           "layout minor_to_major field contains %d elements, "
    188           "but shape is rank %lld: {%s}; shape: %s",
    189           layout.minor_to_major_size(), ShapeUtil::Rank(shape),
    190           tensorflow::str_util::Join(layout.minor_to_major(), ", ").c_str(),
    191           shape.ShortDebugString().c_str());
    192     }
    193 
    194     std::vector<bool> dimensions_in_layout(ShapeUtil::Rank(shape), false);
    195     for (int64 i = 0; i < ShapeUtil::Rank(shape); ++i) {
    196       int64 dim = layout.minor_to_major(i);
    197       if (dim < 0 || dim >= ShapeUtil::Rank(shape)) {
    198         return InvalidArgument(
    199             "layout minor_to_major field has out-of-bounds value: %s",
    200             HumanString(layout).c_str());
    201       }
    202       if (dimensions_in_layout[dim]) {
    203         return InvalidArgument(
    204             "layout minor_to_major field has duplicate values: {%s}",
    205             HumanString(layout).c_str());
    206       }
    207       dimensions_in_layout[dim] = true;
    208     }
    209 
    210     if (layout.padded_dimensions_size() > 0) {
    211       if (layout.padded_dimensions_size() != ShapeUtil::Rank(shape)) {
    212         return InvalidArgument(
    213             "layout has %d padded dimensions, but shape is rank %lld",
    214             layout.padded_dimensions_size(), ShapeUtil::Rank(shape));
    215       }
    216       for (int i = 0; i < layout.padded_dimensions_size(); ++i) {
    217         if (layout.padded_dimensions(i) < shape.dimensions(i)) {
    218           return InvalidArgument(
    219               "for dimension %d, dimension padding (%lld) is smaller than "
    220               "the dimension size (%lld) of the shape",
    221               i, layout.padded_dimensions(i), shape.dimensions(i));
    222         }
    223       }
    224     }
    225   }
    226 
    227   return tensorflow::Status::OK();
    228 }
    229 
    230 /* static */ void LayoutUtil::ClearLayout(Shape* shape) {
    231   shape->clear_layout();
    232   for (auto& element_shape : *shape->mutable_tuple_shapes()) {
    233     ClearLayout(&element_shape);
    234   }
    235 }
    236 
    237 /* static */ void LayoutUtil::ClearLayout(ProgramShape* program_shape) {
    238   for (auto& parameter_shape : *program_shape->mutable_parameters()) {
    239     LayoutUtil::ClearLayout(&parameter_shape);
    240   }
    241   LayoutUtil::ClearLayout(program_shape->mutable_result());
    242 }
    243 
    244 /* static */ bool LayoutUtil::IsDenseArray(const Shape& shape) {
    245   return ShapeUtil::IsArray(shape) && shape.has_layout() &&
    246          IsDense(shape.layout());
    247 }
    248 
    249 /* static */ bool LayoutUtil::IsDense(const Layout& layout) {
    250   return layout.format() == DENSE;
    251 }
    252 
    253 /* static */ bool LayoutUtil::IsMonotonicWithDim0Minor(const Layout& layout) {
    254   CHECK(layout.format() == DENSE);
    255   return std::is_sorted(layout.minor_to_major().begin(),
    256                         layout.minor_to_major().end());
    257 }
    258 
    259 /* static */ bool LayoutUtil::IsMonotonicWithDim0Major(const Layout& layout) {
    260   CHECK(layout.format() == DENSE);
    261   return std::is_sorted(layout.minor_to_major().begin(),
    262                         layout.minor_to_major().end(), std::greater<int64>());
    263 }
    264 
    265 /* static */ bool LayoutUtil::IsPadded(const Shape& shape) {
    266   if (ShapeUtil::IsTuple(shape) || !HasLayout(shape) ||
    267       shape.layout().padded_dimensions_size() == 0) {
    268     return false;
    269   }
    270   CHECK(IsDenseArray(shape));
    271   CHECK_EQ(shape.dimensions_size(), shape.layout().padded_dimensions_size());
    272   for (int64 i = 0; i < shape.dimensions_size(); ++i) {
    273     if (shape.layout().padded_dimensions(i) > shape.dimensions(i)) {
    274       return true;
    275     }
    276   }
    277   return false;
    278 }
    279 
    280 /* static */ tensorflow::gtl::ArraySlice<int64> LayoutUtil::PaddedDimensions(
    281     const Shape& shape) {
    282   CHECK(IsDenseArray(shape));
    283   return AsInt64Slice(shape.layout().padded_dimensions());
    284 }
    285 
    286 /* static */ int64 LayoutUtil::PaddedDimension(const Shape& shape,
    287                                                int64 index) {
    288   CHECK(IsDenseArray(shape));
    289   return shape.layout().padded_dimensions(index);
    290 }
    291 
    292 /* static */ PaddingValue LayoutUtil::GetPaddingValue(const Shape& shape) {
    293   CHECK(IsDenseArray(shape));
    294   return shape.layout().padding_value();
    295 }
    296 
    297 /* static */ bool LayoutUtil::IsSparseArray(const Shape& shape) {
    298   return ShapeUtil::IsArray(shape) && shape.has_layout() &&
    299          IsSparse(shape.layout());
    300 }
    301 
    302 /* static */ bool LayoutUtil::IsSparse(const Layout& layout) {
    303   return layout.format() == SPARSE;
    304 }
    305 
    306 /* static */ int64 LayoutUtil::MaxSparseElements(const Layout& layout) {
    307   CHECK(IsSparse(layout));
    308   return layout.max_sparse_elements();
    309 }
    310 
    311 /* static */ bool LayoutUtil::HasLayout(const Shape& shape) {
    312   if (ShapeUtil::IsTuple(shape)) {
    313     // Tuple shape: all subshapes must have a layout.
    314     return std::all_of(shape.tuple_shapes().begin(), shape.tuple_shapes().end(),
    315                        [](const Shape& s) { return HasLayout(s); });
    316   } else if (ShapeUtil::IsOpaque(shape)) {
    317     return true;
    318   }
    319   return shape.has_layout() && shape.layout().format() != INVALID_FORMAT;
    320 }
    321 
    322 /* static */ bool LayoutUtil::HasLayout(const ProgramShape& program_shape) {
    323   for (auto& parameter_shape : program_shape.parameters()) {
    324     if (!LayoutUtil::HasLayout(parameter_shape)) {
    325       return false;
    326     }
    327   }
    328   return LayoutUtil::HasLayout(program_shape.result());
    329 }
    330 
    331 /* static */ bool LayoutUtil::Equal(const Layout& lhs, const Layout& rhs) {
    332   return protobuf_util::ProtobufEquals(lhs, rhs);
    333 }
    334 
    335 /* static */ tensorflow::gtl::ArraySlice<int64> LayoutUtil::MinorToMajor(
    336     const Shape& shape) {
    337   CHECK(IsDenseArray(shape));
    338   return AsInt64Slice(shape.layout().minor_to_major());
    339 }
    340 
    341 /* static */ tensorflow::gtl::ArraySlice<int64> LayoutUtil::MinorToMajor(
    342     const Layout& layout) {
    343   CHECK(layout.format() == DENSE);
    344   return AsInt64Slice(layout.minor_to_major());
    345 }
    346 
    347 /* static */ int64 LayoutUtil::Major(const Layout& layout,
    348                                      int64 physical_dimension_number) {
    349   CHECK_LE(0, physical_dimension_number);
    350   CHECK_LT(physical_dimension_number, layout.minor_to_major_size());
    351   return Minor(layout,
    352                layout.minor_to_major_size() - 1 - physical_dimension_number);
    353 }
    354 
    355 /* static */ int64 LayoutUtil::Minor(const Layout& layout,
    356                                      int64 physical_dimension_number) {
    357   CHECK_EQ(layout.format(), DENSE);
    358   CHECK_LE(0, physical_dimension_number);
    359   CHECK_LT(physical_dimension_number, layout.minor_to_major_size());
    360   return layout.minor_to_major(physical_dimension_number);
    361 }
    362 
    363 /* static */ std::vector<int64> LayoutUtil::MakeLogicalToPhysical(
    364     const Layout& layout) {
    365   std::vector<int64> logical_to_physical(layout.minor_to_major_size());
    366   for (int64 physical = 0; physical < logical_to_physical.size(); ++physical) {
    367     const int64 logical = Major(layout, physical);
    368     logical_to_physical[logical] = physical;
    369   }
    370   return logical_to_physical;
    371 }
    372 
    373 /* static */ string LayoutUtil::HumanString(const Layout& layout) {
    374   if (IsSparse(layout)) {
    375     return tensorflow::strings::StrCat("sparse{", layout.max_sparse_elements(),
    376                                        "}");
    377   }
    378   CHECK(IsDense(layout));
    379   return tensorflow::strings::StrCat(
    380       "{", tensorflow::str_util::Join(layout.minor_to_major(), ","), "}");
    381 }
    382 
    383 namespace {
    384 
    385 // Internal helper for recursively copying layouts.
    386 tensorflow::Status CopyLayoutInternal(const Shape& src, Shape* dst) {
    387   if (ShapeUtil::IsTuple(src) != ShapeUtil::IsTuple(*dst)) {
    388     return InvalidArgument(
    389         "cannot copy layout from shape: shape structure differs");
    390   }
    391   if (ShapeUtil::IsTuple(src)) {
    392     if (ShapeUtil::TupleElementCount(src) !=
    393         ShapeUtil::TupleElementCount(*dst)) {
    394       return InvalidArgument(
    395           "cannot copy layout from shape: tuple element count differs");
    396     }
    397     for (int64 i = 0; i < ShapeUtil::TupleElementCount(src); ++i) {
    398       TF_RETURN_IF_ERROR(CopyLayoutInternal(src.tuple_shapes(i),
    399                                             dst->mutable_tuple_shapes(i)));
    400     }
    401   } else {
    402     if (src.has_layout()) {
    403       if (ShapeUtil::Rank(src) != ShapeUtil::Rank(*dst)) {
    404         return InvalidArgument("cannot copy layout from shape: ranks differs");
    405       }
    406       TF_RETURN_IF_ERROR(
    407           LayoutUtil::ValidateLayoutForShape(src.layout(), *dst));
    408       *dst->mutable_layout() = src.layout();
    409     } else {
    410       dst->clear_layout();
    411     }
    412   }
    413   return tensorflow::Status::OK();
    414 }
    415 
    416 }  // namespace
    417 
    418 /* static */
    419 tensorflow::Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src,
    420                                                        Shape* dst) {
    421   return CopyLayoutInternal(src, dst);
    422 }
    423 
    424 /* static */ bool LayoutUtil::LayoutsInShapesEqual(const Shape& lhs,
    425                                                    const Shape& rhs) {
    426   if (ShapeUtil::IsTuple(lhs) != ShapeUtil::IsTuple(rhs)) {
    427     return false;
    428   }
    429   if (ShapeUtil::IsTuple(lhs)) {
    430     if (ShapeUtil::TupleElementCount(lhs) !=
    431         ShapeUtil::TupleElementCount(rhs)) {
    432       return false;
    433     }
    434     for (int i = 0; i < ShapeUtil::TupleElementCount(lhs); ++i) {
    435       if (!LayoutsInShapesEqual(lhs.tuple_shapes(i), rhs.tuple_shapes(i))) {
    436         return false;
    437       }
    438     }
    439     return true;
    440   } else {
    441     return ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs) &&
    442            LayoutUtil::Equal(lhs.layout(), rhs.layout());
    443   }
    444 }
    445 
    446 /* static */ bool LayoutUtil::AreDimensionsConsecutive(
    447     const Layout& layout, tensorflow::gtl::ArraySlice<int64> dims) {
    448   CHECK(IsDense(layout));
    449   std::vector<int64> positions_in_layout;
    450   for (int64 dim : dims) {
    451     positions_in_layout.push_back(
    452         PositionInContainer(layout.minor_to_major(), dim));
    453   }
    454   std::sort(positions_in_layout.begin(), positions_in_layout.end());
    455   for (size_t i = 1; i < positions_in_layout.size(); ++i) {
    456     if (1 != positions_in_layout[i] - positions_in_layout[i - 1]) {
    457       return false;
    458     }
    459   }
    460   return true;
    461 }
    462 
    463 std::ostream& operator<<(std::ostream& out, const Layout& layout) {
    464   out << LayoutUtil::HumanString(layout);
    465   return out;
    466 }
    467 
    468 }  // namespace xla
    469