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 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ 18 19 #include <vector> 20 21 #include "tensorflow/compiler/xla/shape_util.h" 22 23 namespace xla { 24 namespace cpu { 25 26 // ShapePartitionAssigner partitions the most-major dimensions of 'shape' such 27 // that the total partition count <= 'target_partition_count'. 28 // 29 // Example 1: 30 // 31 // Let 'shape' = [8, 16, 32] and 'target_partition_count' = 6. 32 // 33 // Because the most-major dimension size is <= 'target_partition_count', we 34 // can generate our target number of partitions by partition the most-major 35 // dimensions. 36 // 37 // This will result in the following partitions of the most-major dimension: 38 // 39 // [0, 1), [1, 2), [2, 3), [3, 4), [4, 5) [5, 8) 40 // 41 // Note that the last partition has residual because the dimension size is 42 // not a multiple of the partition count. 43 // 44 // 45 // Example 2: 46 // 47 // Let 'shape' = [8, 16, 32] and 'target_partition_count' = 16. 48 // 49 // Because the most-major dimension only has size 8, we must also partition 50 // the next most-major dimension to generate the target of 16 partitions. 51 // We factor 'target_partition_count' by the number of most-major dimensions 52 // we need to partition, to get a per-dimension target partition count: 53 // 54 // target_dimension_partition_count = 16 ^ (1 / 2) == 4 55 // 56 // This will result in the following partitions of the most-major dimension: 57 // 58 // [0, 2), [2, 4), [4, 6), [6, 8) 59 // 60 // This will result in the following partitions of the second most-major 61 // dimension: 62 // 63 // [0, 4), [4, 8), [8, 12), [12, 16) 64 // 65 class ShapePartitionAssigner { 66 public: ShapePartitionAssigner(const Shape & shape)67 ShapePartitionAssigner(const Shape& shape) : shape_(shape) {} 68 69 // Returns dimension partition counts (starting at outer-most dimension). 70 std::vector<int64_t> Run(int64_t target_partition_count); 71 72 // Returns the total partition count based on 'dimension_partition_counts'. 73 static int64_t GetTotalPartitionCount( 74 const std::vector<int64_t>& dimension_partition_counts); 75 76 private: 77 const Shape& shape_; 78 }; 79 80 // ShapePartitionIterator iterates through outer-dimension partitions of 81 // 'shape' as specified by 'dimension_partition_counts'. 82 class ShapePartitionIterator { 83 public: 84 ShapePartitionIterator( 85 const Shape& shape, 86 const std::vector<int64_t>& dimension_partition_counts); 87 88 // Returns a partition [start, size] for each dimension. 89 // Partitions are listed starting from outer-most dimension first. 90 std::vector<std::pair<int64_t, int64_t>> GetPartition(int64_t index) const; 91 92 int64_t GetTotalPartitionCount() const; 93 94 private: 95 const Shape& shape_; 96 const std::vector<int64_t> dimension_partition_counts_; 97 98 std::vector<int64_t> dimensions_; 99 std::vector<int64_t> dimension_partition_sizes_; 100 std::vector<int64_t> dimension_partition_strides_; 101 }; 102 103 } // namespace cpu 104 } // namespace xla 105 106 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ 107