Home | History | Annotate | Download | only in service 
      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_COMPUTATION_PLACER_H_
     17 #define TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_PLACER_H_
     18 
     19 #include <map>
     20 #include <memory>
     21 #include <vector>
     22 
     23 #include "tensorflow/compiler/xla/array2d.h"
     24 #include "tensorflow/compiler/xla/status.h"
     25 #include "tensorflow/compiler/xla/statusor.h"
     26 #include "tensorflow/compiler/xla/xla_data.pb.h"
     27 #include "tensorflow/core/lib/core/status.h"
     28 #include "tensorflow/core/platform/macros.h"
     29 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
     30 #include "tensorflow/core/platform/types.h"
     31 
     32 namespace xla {
     33 
     34 // Class that represents the device assignment for a set of XLA replicated
     35 // computations. For R replicas and C computations, R * C devices are required
     36 // execute the computation in parallel. The assigned device ids can be accessed
     37 // by assignment(replica, computation).
     38 class DeviceAssignment : public Array2D<int> {
     39  public:
     40   DeviceAssignment() {}
     41   DeviceAssignment(int replica_count, int computation_count)
     42       : Array2D<int>(replica_count, computation_count, -1) {
     43     CHECK_GT(replica_count, 0);
     44     CHECK_GT(computation_count, 0);
     45   }
     46 
     47   int replica_count() const { return height(); }
     48   int computation_count() const { return width(); }
     49 
     50   // Protocol buffer serialization and deserialization.
     51   Status Serialize(DeviceAssignmentProto* proto) const;
     52 
     53   // Return a std::unique_ptr<DeviceAssignment> instead of a DeviceAssignment
     54   // directly because one of the supported TF platforms (mac) does not compile
     55   // due to a StatusOr of an incomplete type (DeviceAssignment).
     56   static StatusOr<std::unique_ptr<DeviceAssignment>> Deserialize(
     57       const DeviceAssignmentProto& proto);
     58 };
     59 
     60 // A generic implementation of the XLA computation placer, which assigns device
     61 // ids to a set of replicated computations.
     62 class ComputationPlacer {
     63  public:
     64   ComputationPlacer() {}
     65   virtual ~ComputationPlacer() {}
     66 
     67   // Returns the device id assigned to the given replica and computation
     68   // instance for [replica_count x computation_count] setup. The returned device
     69   // id must match the assignement from PlaceReplicatedComputation().
     70   virtual StatusOr<int> DeviceId(int replica, int computation,
     71                                  int replica_count, int computation_count);
     72 
     73   // Returns the device ids assigned to a set of replicated computations, given
     74   // the number of replicas and the number of computations.
     75   virtual StatusOr<DeviceAssignment> AssignDevices(int replica_count,
     76                                                    int computation_count);
     77 
     78   using ComputationPlacerCreationFunction =
     79       std::unique_ptr<ComputationPlacer> (*)();
     80 
     81   // Registers a computation placer creation function for a particular platform.
     82   static void RegisterComputationPlacer(
     83       perftools::gputools::Platform::Id platform_id,
     84       ComputationPlacerCreationFunction creation_function);
     85 
     86   // Returns the computation placer singleton pointer if it is available for the
     87   // given platform, or an error status if it is not.
     88   static StatusOr<ComputationPlacer*> GetForPlatform(
     89       const perftools::gputools::Platform* platform);
     90 
     91  private:
     92   // The mutex that guards the platform-to-computation placer map.
     93   static tensorflow::mutex platform_computation_placer_mutex_;
     94 
     95   // State kept for each kind of ComputationPlacer. Registration functions set
     96   // up creation_function, and then we use that to lazily create "placer" the
     97   // first time GetForPlatform is invoked for a particular id.
     98   struct State {
     99     std::unique_ptr<ComputationPlacer> placer;
    100     ComputationPlacerCreationFunction creation_function = nullptr;
    101   };
    102 
    103   // Map from platform kind to computation placer singleton.
    104   static std::map<perftools::gputools::Platform::Id, State>*
    105   GetPlatformComputationPlacers();
    106 
    107   perftools::gputools::Platform::Id platform_id_;
    108 
    109   TF_DISALLOW_COPY_AND_ASSIGN(ComputationPlacer);
    110 };
    111 
    112 }  // namespace xla
    113 
    114 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_PLACER_H_
    115