Home | History | Annotate | Download | only in interpreter 
      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/service/interpreter/executable.h"
     17 
     18 #include <algorithm>
     19 #include <cstring>
     20 #include <string>
     21 #include <utility>
     22 #include <vector>
     23 
     24 #include "tensorflow/compiler/xla/literal_util.h"
     25 #include "tensorflow/compiler/xla/ptr_util.h"
     26 #include "tensorflow/compiler/xla/service/hlo_computation.h"
     27 #include "tensorflow/compiler/xla/service/hlo_evaluator.h"
     28 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
     29 #include "tensorflow/compiler/xla/service/interpreter/executor.h"
     30 #include "tensorflow/compiler/xla/service/transfer_manager.h"
     31 #include "tensorflow/compiler/xla/shape_util.h"
     32 #include "tensorflow/compiler/xla/status_macros.h"
     33 #include "tensorflow/core/lib/core/errors.h"
     34 #include "tensorflow/core/platform/env.h"
     35 #include "tensorflow/core/platform/mutex.h"
     36 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
     37 
     38 namespace xla {
     39 namespace interpreter {
     40 
     41 namespace se = ::perftools::gputools;
     42 
     43 InterpreterExecutable::InterpreterExecutable(
     44     std::unique_ptr<const HloModule> hlo_module)
     45     : Executable(std::move(hlo_module), /*hlo_profile_printer=*/nullptr,
     46                  /*hlo_profile_index_map=*/nullptr) {}
     47 
     48 InterpreterExecutable::~InterpreterExecutable() {}
     49 
     50 StatusOr<std::unique_ptr<ShapedBuffer>> InterpreterExecutable::ExecuteOnStream(
     51     const ServiceExecutableRunOptions* run_options,
     52     tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments,
     53     HloExecutionProfile* hlo_execution_profile) {
     54   se::Stream* stream = run_options->stream();
     55   se::StreamExecutor* executor = stream->parent();
     56   const se::Platform* platform = executor->platform();
     57 
     58   VLOG(1) << "Execute " << module().name();
     59   if (VLOG_IS_ON(2)) {
     60     for (const auto& a : arguments) {
     61       VLOG(2) << "-- argument " << *a;
     62     }
     63   }
     64 
     65   uint64 start_micros = tensorflow::Env::Default()->NowMicros();
     66 
     67   const HloComputation* computation = module().entry_computation();
     68   if (computation->num_parameters() != arguments.size()) {
     69     return tensorflow::errors::Internal(
     70         "Mismatch between argument count and graph parameter count.");
     71   }
     72 
     73   TF_ASSIGN_OR_RETURN(TransferManager * transfer_manager,
     74                       TransferManager::GetForPlatform(platform));
     75 
     76   // Transform the ShapedBuffer arguments into literals which the evaluator
     77   // consumes.
     78   std::vector<std::unique_ptr<Literal>> arg_literals;
     79   for (int64 p = 0; p < computation->num_parameters(); ++p) {
     80     TF_ASSIGN_OR_RETURN(
     81         std::unique_ptr<Literal> arg_literal,
     82         transfer_manager->TransferLiteralFromDevice(executor, *arguments[p]));
     83     arg_literals.push_back(std::move(arg_literal));
     84   }
     85 
     86   // Execute the graph using the HloEvaluator.
     87   HloEvaluator evaluator;
     88   TF_ASSIGN_OR_RETURN(
     89       std::unique_ptr<Literal> result_literal,
     90       evaluator.Evaluate<std::unique_ptr<Literal>>(*computation, arg_literals));
     91 
     92   // Transform the result literal back into a ShapedBuffer.
     93   TF_ASSIGN_OR_RETURN(std::unique_ptr<ShapedBuffer> result,
     94                       transfer_manager->AllocateShapedBuffer(
     95                           result_literal->shape(), run_options->allocator(),
     96                           run_options->device_ordinal()));
     97   TF_RETURN_IF_ERROR(transfer_manager->TransferLiteralToDevice(
     98       executor, *result_literal, *result));
     99 
    100   uint64 end_micros = tensorflow::Env::Default()->NowMicros();
    101 
    102   {
    103     tensorflow::mutex_lock lock(mutex_);
    104     const double nanoseconds = (end_micros - start_micros) * 1000.0;
    105     execution_profile_.set_compute_time_ns(std::max(nanoseconds, 1.0));
    106   }
    107 
    108   return std::move(result);
    109 }
    110 
    111 StatusOr<std::unique_ptr<ShapedBuffer>>
    112 InterpreterExecutable::ExecuteAsyncOnStream(
    113     const ServiceExecutableRunOptions* run_options,
    114     tensorflow::gtl::ArraySlice<const ShapedBuffer*> arguments) {
    115   return tensorflow::errors::Unimplemented(
    116       "ExecuteAsyncOnStream is not yet supported on Interpreter.");
    117 }
    118 
    119 /*static*/ int64 InterpreterExecutable::ShapeSizeBytes(const Shape& shape) {
    120   if (ShapeUtil::IsOpaque(shape)) {
    121     return sizeof(void*);
    122   }
    123   return ShapeUtil::ByteSizeOf(shape, sizeof(void*));
    124 }
    125 
    126 }  // namespace interpreter
    127 }  // namespace xla
    128