<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/jit/xla_device_context.h"

#include <memory>

#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/tf2xla/literal_util.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/framework/tensor_reference.h"
#include "tensorflow/core/platform/mem.h"
#include "tensorflow/stream_executor/platform/port.h"

namespace tensorflow {

// The allocator used for Tensors assigned to the XLA device.
XlaDeviceAllocator::XlaDeviceAllocator(
    stream_executor::StreamExecutor* stream_executor)
    : stream_executor_(stream_executor) {}

XlaDeviceAllocator::~XlaDeviceAllocator() = default;

string XlaDeviceAllocator::Name() { return "xla"; }

void* XlaDeviceAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
  // We always return an empty XlaTensor object, encoded as an opaque tagged
  // pointer. We can return an empty object and ignore num_bytes here because we
  // have control over all of the uses of this device tensor, and can lazily
  // allocate memory when used. This allows us to also know the shape of the
  // allocated Tensor, which is useful if the device's tensor representation
  // differs from the host.
  return XlaTensor::ToOpaquePointer(new XlaTensor());
}

void XlaDeviceAllocator::DeallocateRaw(void* ptr) {
  delete XlaTensor::FromOpaquePointer(ptr);
}

absl::optional<AllocatorStats> XlaDeviceAllocator::GetStats() {
  absl::optional<stream_executor::AllocatorStats> se_stats =
      stream_executor_->GetAllocatorStats();
  if (!se_stats) {
    return absl::nullopt;
  }

  tensorflow::AllocatorStats tf_stats;
  tf_stats.num_allocs = se_stats->num_allocs;
  tf_stats.bytes_in_use = se_stats->bytes_in_use;
  tf_stats.peak_bytes_in_use = se_stats->peak_bytes_in_use;
  tf_stats.largest_alloc_size = se_stats->largest_alloc_size;
  tf_stats.bytes_limit = se_stats->bytes_limit;
  tf_stats.bytes_reserved = se_stats->bytes_reserved;
  tf_stats.peak_bytes_reserved = se_stats->peak_bytes_reserved;
  tf_stats.bytes_reservable_limit = se_stats->bytes_reservable_limit;
  return tf_stats;
}

XlaDeviceContext::XlaDeviceContext(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream,
    std::vector<std::shared_ptr<se::Stream>> device_to_device_streams,
    xla::LocalClient* client,
    XlaCompiler::ShapeRepresentationFn shape_representation_fn,
    thread::ThreadPool* thread_pool, bool use_fast_mem)
    : stream_(std::move(compute_stream)),
      host_to_device_stream_(std::move(host_to_device_stream)),
      device_to_host_stream_(std::move(device_to_host_stream)),
      device_to_device_streams_(std::move(device_to_device_streams)),
      client_(client),
      transfer_manager_(client->backend().transfer_manager()),
      shape_representation_fn_(std::move(shape_representation_fn)),
      thread_pool_(thread_pool),
      use_fast_mem_(use_fast_mem) {
  CHECK(host_to_device_stream_ != nullptr);
  CHECK(stream_ != nullptr);
  if (!shape_representation_fn_) {
    shape_representation_fn_ =
        [](const TensorShape& shape, DataType dtype,
           bool use_fast_memory) -> xla::StatusOr<xla::Shape> {
      xla::Shape xla_shape;
      TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape));
      return xla_shape;
    };
  }
}

void XlaDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,
                                              Device* device,
                                              Tensor* output_tensor,
                                              StatusCallback done) const {
  done(errors::Unimplemented("XLA->XLA same-device copies not implemented."));
}

void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                             Device* device,
                                             Tensor* device_tensor,
                                             StatusCallback done,
                                             bool sync_dst_compute) const {
  if (cpu_tensor->NumElements() == 0) {
    VLOG(2) << "CopyCPUTensorToDevice empty tensor";
    done(Status::OK());
    return;
  }

  VLOG(2) << "CopyCPUTensorToDevice use_fast_mem " << use_fast_mem_ << " "
          << this << " "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " " << cpu_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  CHECK(xla_tensor);

  Status status = [&]() -> Status {
    TF_ASSIGN_OR_RETURN(
        xla::Shape shape,
        shape_representation_fn_(device_tensor->shape(), device_tensor->dtype(),
                                 use_fast_mem_));

    // The device tensor should always be fresh.
    TF_RET_CHECK(!xla_tensor->has_shaped_buffer());

    TF_RETURN_IF_ERROR(
        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal()));

    // The cpu_tensor and literal that we created here hold the data of host
    // tensor in descending layout. The layout could be different from layout in
    // device_tensor (but the logical shape has to be the same). The
    // transfer_manager is responsible to do corresponding transposing when
    // transferring the data to device.
    xla::BorrowingLiteral literal(
        static_cast<const char*>(DMAHelper::base(cpu_tensor)),
        xla::ShapeUtil::MakeShape(shape.element_type(),
                                  xla::AsInt64Slice(shape.dimensions())));

    VLOG(2) << "Transfer to device as literal: " << literal.ToString() << " "
            << xla_tensor->shaped_buffer().ToString();
    if (UseMultipleStreams() &&
        !transfer_manager_->CanShapedBufferBeAccessedNow(
            stream_->parent(), xla_tensor->shaped_buffer())) {
      // Initially wait for the compute stream so that memory allocations are
      // synchronized.
      host_to_device_stream_->ThenWaitFor(stream_.get());
    }

    TF_RETURN_IF_ERROR(transfer_manager_->TransferLiteralToDeviceAsync(
        host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer()));

    if (UseMultipleStreams()) {
      auto event = std::make_shared<se::Event>(stream_->parent());
      TF_RET_CHECK(event->Init()) << "Event failed to initialize!";
      host_to_device_stream_->ThenRecordEvent(event.get());
      xla_tensor->ResetDefinitionEvent(std::move(event),
                                       host_to_device_stream_.get());
    }

    return Status::OK();
  }();
  if (!status.ok()) {
    done(status);
    return;
  }

  // Create a reference to hold onto cpu_tensor until after the literal has
  // been transferred
  TensorReference ref(*cpu_tensor);
  if (UseMultipleStreams()) {
    // Unref the host tensor when the transfer completes.
    // We don't defer the call to done() onto the stream here, and the reasons
    // why this is correct are subtle. We assume that:
    // a) all consumers of the device tensor will wait for its definition event.
    // b) if the tensor is destroyed, then the memory allocator will not hand
    //    out the same buffers until the transfer has completed.
    host_to_device_stream_->ThenDoHostCallback([ref]() { ref.Unref(); });
    done(status);
  } else {
    host_to_device_stream_->ThenDoHostCallback([ref, done]() {
      ref.Unref();
      done(Status::OK());
    });
  }
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             absl::string_view tensor_name,
                                             Device* device, Tensor* cpu_tensor,
                                             StatusCallback done) {
  if (device_tensor->NumElements() == 0) {
    VLOG(2) << "CopyDeviceTensorToCPU empty tensor";
    done(Status::OK());
    return;
  }
  VLOG(2) << "CopyDeviceTensorToCPU "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " " << device_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  std::shared_ptr<se::Stream> device_to_host_stream;
  if (device_to_host_stream_) {
    device_to_host_stream = device_to_host_stream_;
  } else {
    stream_executor::port::StatusOr<xla::StreamPool::Ptr> ptr_or_status =
        client_->mutable_backend()->BorrowStream(
            stream_->parent()->device_ordinal());
    if (!ptr_or_status.status().ok()) {
      done(ptr_or_status.status());
      return;
    }
    device_to_host_stream =
        std::shared_ptr<se::Stream>(std::move(ptr_or_status.ValueOrDie()));
  }

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream.get());

  // Transfer manager requires the shape of the shaped buffer to be the same as
  // literal shape except for the layout.  Set the literal to use xla_tensor's
  // shape as it is derived from the cpu_tensor's shape using
  // shape_representation_fn_.
  xla::MutableBorrowingLiteral literal;
  TF_CHECK_OK(HostTensorToMutableBorrowingLiteral(
      xla::LayoutUtil::GetWithDefaultLayout(
          xla_tensor->shaped_buffer().on_host_shape()),
      cpu_tensor, &literal));

  TensorReference ref(*device_tensor);
  const bool device_allows_sync_on_completion =
      device->AllowsSyncOnCompletion();
  // Explicitly capture device_to_host_stream to make sure the stream is alive
  // before the transfer finishes.
  transfer_manager_->TransferLiteralFromDevice(
      device_to_host_stream.get(), xla_tensor->shaped_buffer(), literal,
      [ref, xla_tensor, done, device_to_host_stream,
       device_allows_sync_on_completion](xla::Status status) {
        Status done_status = status;
        VLOG(2) << "Transfer from device as literal: "
                << xla_tensor->shaped_buffer().ToString();
        // For devices don't allow sync on completion, the device execution is
        // deferred. We check the execution stream status here to avoid wrong
        // results from a failed stream being propagated to following
        // host-side ops.
        if (!device_allows_sync_on_completion) {
          done_status.Update(xla_tensor->RefreshStatusOfStreams());
        }
        done(done_status);
        ref.Unref();
      });
}

se::Stream* XlaDeviceContext::GetDeviceToDeviceStream() {
  DCHECK_GT(device_to_device_streams_.size(), 0);
  absl::MutexLock lock(&mu_);
  int stream = next_stream_;
  next_stream_ = (next_stream_ + 1) % device_to_device_streams_.size();
  return device_to_device_stream(stream);
}

}  // namespace tensorflow
<commit_msg>Avoid deleting a bad Stream in a host callback enqueued on itself, which can cause bad behaviors on some platforms.<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/jit/xla_device_context.h"

#include <memory>

#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/tf2xla/literal_util.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/framework/tensor_reference.h"
#include "tensorflow/core/platform/mem.h"
#include "tensorflow/stream_executor/platform/port.h"

namespace tensorflow {

// The allocator used for Tensors assigned to the XLA device.
XlaDeviceAllocator::XlaDeviceAllocator(
    stream_executor::StreamExecutor* stream_executor)
    : stream_executor_(stream_executor) {}

XlaDeviceAllocator::~XlaDeviceAllocator() = default;

string XlaDeviceAllocator::Name() { return "xla"; }

void* XlaDeviceAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
  // We always return an empty XlaTensor object, encoded as an opaque tagged
  // pointer. We can return an empty object and ignore num_bytes here because we
  // have control over all of the uses of this device tensor, and can lazily
  // allocate memory when used. This allows us to also know the shape of the
  // allocated Tensor, which is useful if the device's tensor representation
  // differs from the host.
  return XlaTensor::ToOpaquePointer(new XlaTensor());
}

void XlaDeviceAllocator::DeallocateRaw(void* ptr) {
  delete XlaTensor::FromOpaquePointer(ptr);
}

absl::optional<AllocatorStats> XlaDeviceAllocator::GetStats() {
  absl::optional<stream_executor::AllocatorStats> se_stats =
      stream_executor_->GetAllocatorStats();
  if (!se_stats) {
    return absl::nullopt;
  }

  tensorflow::AllocatorStats tf_stats;
  tf_stats.num_allocs = se_stats->num_allocs;
  tf_stats.bytes_in_use = se_stats->bytes_in_use;
  tf_stats.peak_bytes_in_use = se_stats->peak_bytes_in_use;
  tf_stats.largest_alloc_size = se_stats->largest_alloc_size;
  tf_stats.bytes_limit = se_stats->bytes_limit;
  tf_stats.bytes_reserved = se_stats->bytes_reserved;
  tf_stats.peak_bytes_reserved = se_stats->peak_bytes_reserved;
  tf_stats.bytes_reservable_limit = se_stats->bytes_reservable_limit;
  return tf_stats;
}

XlaDeviceContext::XlaDeviceContext(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream,
    std::vector<std::shared_ptr<se::Stream>> device_to_device_streams,
    xla::LocalClient* client,
    XlaCompiler::ShapeRepresentationFn shape_representation_fn,
    thread::ThreadPool* thread_pool, bool use_fast_mem)
    : stream_(std::move(compute_stream)),
      host_to_device_stream_(std::move(host_to_device_stream)),
      device_to_host_stream_(std::move(device_to_host_stream)),
      device_to_device_streams_(std::move(device_to_device_streams)),
      client_(client),
      transfer_manager_(client->backend().transfer_manager()),
      shape_representation_fn_(std::move(shape_representation_fn)),
      thread_pool_(thread_pool),
      use_fast_mem_(use_fast_mem) {
  CHECK(host_to_device_stream_ != nullptr);
  CHECK(stream_ != nullptr);
  if (!shape_representation_fn_) {
    shape_representation_fn_ =
        [](const TensorShape& shape, DataType dtype,
           bool use_fast_memory) -> xla::StatusOr<xla::Shape> {
      xla::Shape xla_shape;
      TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape));
      return xla_shape;
    };
  }
}

void XlaDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor,
                                              Device* device,
                                              Tensor* output_tensor,
                                              StatusCallback done) const {
  done(errors::Unimplemented("XLA->XLA same-device copies not implemented."));
}

void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                             Device* device,
                                             Tensor* device_tensor,
                                             StatusCallback done,
                                             bool sync_dst_compute) const {
  if (cpu_tensor->NumElements() == 0) {
    VLOG(2) << "CopyCPUTensorToDevice empty tensor";
    done(Status::OK());
    return;
  }

  VLOG(2) << "CopyCPUTensorToDevice use_fast_mem " << use_fast_mem_ << " "
          << this << " "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " " << cpu_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  CHECK(xla_tensor);

  Status status = [&]() -> Status {
    TF_ASSIGN_OR_RETURN(
        xla::Shape shape,
        shape_representation_fn_(device_tensor->shape(), device_tensor->dtype(),
                                 use_fast_mem_));

    // The device tensor should always be fresh.
    TF_RET_CHECK(!xla_tensor->has_shaped_buffer());

    TF_RETURN_IF_ERROR(
        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal()));

    // The cpu_tensor and literal that we created here hold the data of host
    // tensor in descending layout. The layout could be different from layout in
    // device_tensor (but the logical shape has to be the same). The
    // transfer_manager is responsible to do corresponding transposing when
    // transferring the data to device.
    xla::BorrowingLiteral literal(
        static_cast<const char*>(DMAHelper::base(cpu_tensor)),
        xla::ShapeUtil::MakeShape(shape.element_type(),
                                  xla::AsInt64Slice(shape.dimensions())));

    VLOG(2) << "Transfer to device as literal: " << literal.ToString() << " "
            << xla_tensor->shaped_buffer().ToString();
    if (UseMultipleStreams() &&
        !transfer_manager_->CanShapedBufferBeAccessedNow(
            stream_->parent(), xla_tensor->shaped_buffer())) {
      // Initially wait for the compute stream so that memory allocations are
      // synchronized.
      host_to_device_stream_->ThenWaitFor(stream_.get());
    }

    TF_RETURN_IF_ERROR(transfer_manager_->TransferLiteralToDeviceAsync(
        host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer()));

    if (UseMultipleStreams()) {
      auto event = std::make_shared<se::Event>(stream_->parent());
      TF_RET_CHECK(event->Init()) << "Event failed to initialize!";
      host_to_device_stream_->ThenRecordEvent(event.get());
      xla_tensor->ResetDefinitionEvent(std::move(event),
                                       host_to_device_stream_.get());
    }

    return Status::OK();
  }();
  if (!status.ok()) {
    done(status);
    return;
  }

  // Create a reference to hold onto cpu_tensor until after the literal has
  // been transferred
  TensorReference ref(*cpu_tensor);
  if (UseMultipleStreams()) {
    // Unref the host tensor when the transfer completes.
    // We don't defer the call to done() onto the stream here, and the reasons
    // why this is correct are subtle. We assume that:
    // a) all consumers of the device tensor will wait for its definition event.
    // b) if the tensor is destroyed, then the memory allocator will not hand
    //    out the same buffers until the transfer has completed.
    host_to_device_stream_->ThenDoHostCallback([ref]() { ref.Unref(); });
    done(status);
  } else {
    host_to_device_stream_->ThenDoHostCallback([ref, done]() {
      ref.Unref();
      done(Status::OK());
    });
  }
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             absl::string_view tensor_name,
                                             Device* device, Tensor* cpu_tensor,
                                             StatusCallback done) {
  if (device_tensor->NumElements() == 0) {
    VLOG(2) << "CopyDeviceTensorToCPU empty tensor";
    done(Status::OK());
    return;
  }
  VLOG(2) << "CopyDeviceTensorToCPU "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " " << device_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  std::shared_ptr<se::Stream> device_to_host_stream;
  if (device_to_host_stream_) {
    device_to_host_stream = device_to_host_stream_;
  } else {
    stream_executor::port::StatusOr<xla::StreamPool::Ptr> ptr_or_status =
        client_->mutable_backend()->BorrowStream(
            stream_->parent()->device_ordinal());
    if (!ptr_or_status.status().ok()) {
      done(ptr_or_status.status());
      return;
    }
    device_to_host_stream =
        std::shared_ptr<se::Stream>(std::move(ptr_or_status.ValueOrDie()));
  }

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream.get());

  // Transfer manager requires the shape of the shaped buffer to be the same as
  // literal shape except for the layout.  Set the literal to use xla_tensor's
  // shape as it is derived from the cpu_tensor's shape using
  // shape_representation_fn_.
  xla::MutableBorrowingLiteral literal;
  TF_CHECK_OK(HostTensorToMutableBorrowingLiteral(
      xla::LayoutUtil::GetWithDefaultLayout(
          xla_tensor->shaped_buffer().on_host_shape()),
      cpu_tensor, &literal));

  TensorReference ref(*device_tensor);
  const bool device_allows_sync_on_completion =
      device->AllowsSyncOnCompletion();
  // Explicitly capture device_to_host_stream to make sure the stream is alive
  // before the transfer finishes.
  transfer_manager_->TransferLiteralFromDevice(
      device_to_host_stream.get(), xla_tensor->shaped_buffer(), literal,
      [this, ref, xla_tensor, done, device_to_host_stream,
       device_allows_sync_on_completion](xla::Status status) {
        Status done_status = status;
        VLOG(2) << "Transfer from device as literal: "
                << xla_tensor->shaped_buffer().ToString();
        // For devices don't allow sync on completion, the device execution is
        // deferred. We check the execution stream status here to avoid wrong
        // results from a failed stream being propagated to following
        // host-side ops.
        if (!device_allows_sync_on_completion) {
          done_status.Update(xla_tensor->RefreshStatusOfStreams());
        }
        done(done_status);
        ref.Unref();
        // If a stream is in a bad state, it gets deleted when it's returned to
        // the stream pool, i.e. when it leaves this scope. However, a stream
        // deleting itself in a host callback on itself can cause bad behaviors
        // on some platforms. Releasing it in another stream to avoid that.
        if (!device_allows_sync_on_completion &&
            !device_to_host_stream->RefreshStatus().ok()) {
          auto status_or_new_stream = client_->mutable_backend()->BorrowStream(
              stream_->parent()->device_ordinal());
          if (status_or_new_stream.ok()) {
            status_or_new_stream.ValueOrDie()->ThenDoHostCallback(
                [device_to_host_stream] {});
          }
        }
      });
}

se::Stream* XlaDeviceContext::GetDeviceToDeviceStream() {
  DCHECK_GT(device_to_device_streams_.size(), 0);
  absl::MutexLock lock(&mu_);
  int stream = next_stream_;
  next_stream_ = (next_stream_ + 1) % device_to_device_streams_.size();
  return device_to_device_stream(stream);
}

}  // namespace tensorflow
<|endoftext|>


<commit_before>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/jit/xla_device_context.h"

#include <memory>

#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/tf2xla/literal_util.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/platform/mem.h"

namespace tensorflow {

// The allocator used for Tensors assigned to the XLA device.
XlaDeviceAllocator::XlaDeviceAllocator() {}
XlaDeviceAllocator::~XlaDeviceAllocator() = default;

string XlaDeviceAllocator::Name() { return "xla"; }

void* XlaDeviceAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
  // We always return an empty XlaTensor object, encoded as an opaque tagged
  // pointer. We can return an empty object and ignore num_bytes here because we
  // have control over all of the uses of this device tensor, and can lazily
  // allocate memory when used. This allows us to also know the shape of the
  // allocated Tensor, which is useful if the device's tensor representation
  // differs from the host.
  return XlaTensor::ToOpaquePointer(new XlaTensor());
}

void XlaDeviceAllocator::DeallocateRaw(void* ptr) {
  delete XlaTensor::FromOpaquePointer(ptr);
}

void XlaDeviceAllocator::GetStats(AllocatorStats* stats) { stats->Clear(); }

XlaTransferManager::XlaTransferManager(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream, xla::LocalClient* client,
    bool transfer_as_literal,
    XlaCompiler::ShapeRepresentationFn shape_representation_fn,
    thread::ThreadPool* thread_pool)
    : stream_(std::move(compute_stream)),
      host_to_device_stream_(std::move(host_to_device_stream)),
      device_to_host_stream_(std::move(device_to_host_stream)),
      client_(client),
      transfer_manager_(client->backend().transfer_manager()),
      transfer_as_literal_(transfer_as_literal),
      shape_representation_fn_(std::move(shape_representation_fn)),
      thread_pool_(thread_pool) {
  CHECK(host_to_device_stream_ != nullptr);
  CHECK(device_to_host_stream_ != nullptr);
  CHECK(stream_ != nullptr);
  if (!shape_representation_fn_) {
    shape_representation_fn_ =
        [](const TensorShape& shape,
           DataType dtype) -> xla::StatusOr<TensorShape> { return shape; };
  }
}

Status XlaTransferManager::TransferLiteralToDevice(
    const Tensor& host_tensor, Tensor* device_tensor) const {
  xla::Shape xla_shape;
  TF_RETURN_IF_ERROR(TensorShapeToXLAShape(host_tensor.dtype(),
                                           host_tensor.shape(), &xla_shape));
  // Create a reference to hold onto host_tensor until after the literal has
  // been transferred. Also make sure the literal exists until the function
  // asynchronously completes, as it will be wrapped in an xla::LiteralSlice.
  TensorReference ref(host_tensor);
  auto literal = std::make_shared<xla::BorrowingLiteral>(
      static_cast<const char*>(DMAHelper::base(&host_tensor)), xla_shape);

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer();
  VLOG(1) << "Transfer to device as literal: " << literal->ToString() << " "
          << shaped_buffer.ToString();
  if (UseMultipleStreams() && !transfer_manager_->CanShapedBufferBeAccessedNow(
                                  stream_->parent(), shaped_buffer)) {
    // Initially wait for the compute stream so that memory allocations are
    // synchronized.
    host_to_device_stream_->ThenWaitFor(stream_.get());
  }
  TF_RETURN_IF_ERROR(transfer_manager_->TransferLiteralToDeviceAsync(
      host_to_device_stream_.get(), *literal, shaped_buffer));
  if (UseMultipleStreams()) {
    auto event = std::make_shared<se::Event>(stream_->parent());
    TF_RET_CHECK(event->Init()) << "Event failed to initialize!";
    host_to_device_stream_->ThenRecordEvent(event.get());
    xla_tensor->SetDefinedOn(host_to_device_stream_.get(), std::move(event));
  }
  // Unref the host tensor, and capture the literal shared_ptr too so it goes
  // out of scope when the lambda completes.
  host_to_device_stream_->ThenDoHostCallback([ref, literal]() { ref.Unref(); });

  return Status::OK();
}

void XlaTransferManager::TransferLiteralFromDevice(
    Tensor* host_tensor, const Tensor& device_tensor,
    const StatusCallback& done) const {
  xla::MutableBorrowingLiteral literal;
  TF_CHECK_OK(HostTensorToMutableBorrowingLiteral(host_tensor, &literal));

  const xla::ShapedBuffer& shaped_buffer =
      XlaTensor::FromTensor(&device_tensor)->shaped_buffer();

  TensorReference ref(device_tensor);
  transfer_manager_->TransferLiteralFromDevice(
      device_to_host_stream_.get(), shaped_buffer, literal,
      [=, &shaped_buffer, &literal](xla::Status status) {
        ref.Unref();
        done([&]() -> Status {
          VLOG(1) << "Transfer from device as literal: " << literal.ToString()
                  << " " << shaped_buffer.ToString();
          return status;
        }());
      });
}

void XlaTransferManager::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                               Device* device,
                                               Tensor* device_tensor,
                                               StatusCallback done) const {
  if (cpu_tensor->NumElements() == 0) {
    VLOG(2) << "CopyCPUTensorToDevice empty tensor";
    done(Status::OK());
    return;
  }

  VLOG(2) << "CopyCPUTensorToDevice "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " " << cpu_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  void* src_ptr = const_cast<void*>(DMAHelper::base(cpu_tensor));
  const int64 total_bytes = cpu_tensor->TotalBytes();

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  CHECK(xla_tensor);

  xla::StatusOr<TensorShape> shape_or_status =
      shape_representation_fn_(device_tensor->shape(), device_tensor->dtype());
  if (!shape_or_status.ok()) {
    done(shape_or_status.status());
    return;
  }
  TensorShape shape = shape_or_status.ValueOrDie();
  if (!xla_tensor->has_shaped_buffer()) {
    Status s =
        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal());
    if (!s.ok()) {
      done(s);
      return;
    }
  }

  Status status;
  if (transfer_as_literal_) {
    Tensor reshaped_cpu_tensor;
    if (!reshaped_cpu_tensor.CopyFrom(*cpu_tensor, shape)) {
      done(errors::Internal(
          "Tensor::CopyFrom failed when copying from CPU to XLA device"));
      return;
    }
    status = TransferLiteralToDevice(reshaped_cpu_tensor, device_tensor);
    if (status.ok()) {
      xla_tensor->set_host_tensor(*cpu_tensor);
      host_to_device_stream_->ThenDoHostCallback([this, done]() {
        // We must not call the done closure directly from DoHostCallback
        // to avoid a deadlock. If done() is the callback that ends an
        // Executor's run, the Executor may call XlaDevice::Sync() inside the
        // callback. This deadlocks, because XlaDevice::Sync() waits for all
        // stream activity to complete.
        thread_pool_->Schedule([done]() { done(Status::OK()); });
      });
      return;
    }
  } else {
    se::DeviceMemoryBase dev_dst_ptr =
        XlaTensor::DeviceMemoryFromTensor(*device_tensor);
    host_to_device_stream_->ThenMemcpy(&dev_dst_ptr, src_ptr, total_bytes);
    // TODO(hpucha): Make this asynchronous.
    Status block_status = host_to_device_stream_->BlockHostUntilDone();
    if (!block_status.ok()) {
      status = xla::InternalError(
          "Failed to complete data transfer on stream %p: %s",
          host_to_device_stream_.get(), block_status.error_message().c_str());
    }
  }
  xla_tensor->set_host_tensor(*cpu_tensor);

  done(status);
}

void XlaTransferManager::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                               StringPiece tensor_name,
                                               Device* device,
                                               Tensor* cpu_tensor,
                                               StatusCallback done) {
  if (device_tensor->NumElements() == 0) {
    VLOG(2) << "CopyDeviceTensorToCPU empty tensor";
    done(Status::OK());
    return;
  }
  VLOG(2) << "CopyDeviceTensorToCPU "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " " << device_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  const int64 total_bytes = cpu_tensor->TotalBytes();
  se::DeviceMemoryBase dev_src_ptr =
      XlaTensor::DeviceMemoryFromTensor(*device_tensor);
  void* dst_ptr = DMAHelper::base(cpu_tensor);
  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);

  if (se::Event* event =
          xla_tensor->GetDefinitionEvent(device_to_host_stream_.get())) {
    device_to_host_stream_->ThenWaitFor(event);
    xla_tensor->SetDefinedOn(device_to_host_stream_.get());
  }

  Status status;
  if (transfer_as_literal_) {
    TransferLiteralFromDevice(cpu_tensor, *device_tensor, done);
    return;
  } else {
    device_to_host_stream_->ThenMemcpy(dst_ptr, dev_src_ptr, total_bytes);
    // TODO(hpucha): Make this asynchronous.
    Status block_status = device_to_host_stream_->BlockHostUntilDone();
    if (!block_status.ok()) {
      status = xla::InternalError(
          "Failed to complete data transfer on stream %p: %s", stream_.get(),
          block_status.error_message().c_str());
    }
  }

  done(status);
}

void XlaTransferManager::CopyDeviceTensorToDevice(const Tensor& src_tensor,
                                                  Tensor* dst_tensor,
                                                  const StatusCallback& done) {
  VLOG(2) << "CopyDeviceTensorToDevice "
          << reinterpret_cast<const void*>(src_tensor.tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(dst_tensor->tensor_data().data());
  // Perform memory allocation now, and enqueue the device-to-device transfer.
  Status status = [&]() -> Status {
    if (src_tensor.NumElements() == 0) {
      return Status::OK();
    }
    // TODO(jmolloy): We co-opt the device_to_host stream for device to device
    // transfers; perhaps we should have a dedicated device to device stream? or
    // one per device?
    auto device_to_device_stream = stream_;
    XlaTensor* xla_src = XlaTensor::FromTensor(&src_tensor);
    XlaTensor* xla_dst = XlaTensor::FromTensor(dst_tensor);
    CHECK(xla_src && xla_dst)
        << "Missing destination tensor for device-to-device copy";
    if (!xla_dst->has_shaped_buffer()) {
      TF_ASSIGN_OR_RETURN(
          TensorShape shape,
          shape_representation_fn_(src_tensor.shape(), src_tensor.dtype()));
      TF_RETURN_IF_ERROR(
          xla_dst->AllocateShapedBuffer(src_tensor.dtype(), shape, client_,
                                        stream_->parent()->device_ordinal()));
      if (stream_ != device_to_device_stream) {
        // Initially wait for the compute stream so that memory allocations are
        // synchronized.
        device_to_device_stream->ThenWaitFor(stream_.get());
      }
    }

    if (se::Event* event =
            xla_src->GetDefinitionEvent(device_to_device_stream.get())) {
      device_to_device_stream->ThenWaitFor(event);
      xla_src->SetDefinedOn(device_to_device_stream.get());
    }

    auto from_iter = xla_src->shaped_buffer().buffers().begin();
    auto to_iter = xla_dst->shaped_buffer().buffers().begin();
    for (auto end_iter = xla_src->shaped_buffer().buffers().end();
         from_iter != end_iter; ++from_iter, ++to_iter) {
      device_to_device_stream->ThenMemcpyD2D(
          &to_iter->second, from_iter->second, to_iter->second.size());
    }

    if (UseMultipleStreams()) {
      auto event = std::make_shared<se::Event>(stream_->parent());
      TF_RET_CHECK(event->Init()) << "Event failed to initialize";
      device_to_device_stream->ThenRecordEvent(event.get());
      xla_dst->SetDefinedOn(device_to_device_stream.get(), std::move(event));
    }
    return Status::OK();
  }();
  if (!status.ok()) {
    return done(status);
  } else {
    stream_->ThenDoHostCallback([this, done]() {
      // We must not call the done closure directly from DoHostCallback to avoid
      // a deadlock. If done() is the callback that ends an Executor's run, the
      // Executor may call XlaDevice::Sync() inside the callback. This
      // deadlocks, because XlaDevice::Sync() waits for all stream activity to
      // complete.
      thread_pool_->Schedule([done]() { done(Status::OK()); });
    });
  }
}

XlaDeviceContext::XlaDeviceContext(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream, xla::LocalClient* client,
    bool transfer_as_literal,
    XlaCompiler::ShapeRepresentationFn shape_representation_fn,
    thread::ThreadPool* thread_pool)
    : manager_(std::move(compute_stream), std::move(host_to_device_stream),
               std::move(device_to_host_stream), client, transfer_as_literal,
               std::move(shape_representation_fn), thread_pool) {}

void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                             Device* device,
                                             Tensor* device_tensor,
                                             StatusCallback done) const {
  manager_.CopyCPUTensorToDevice(cpu_tensor, device, device_tensor, done);
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             StringPiece tensor_name,
                                             Device* device, Tensor* cpu_tensor,
                                             StatusCallback done) {
  manager_.CopyDeviceTensorToCPU(device_tensor, tensor_name, device, cpu_tensor,
                                 done);
}

void XlaDeviceContext::CopyDeviceTensorToDevice(const Tensor& src_tensor,
                                                Tensor* dst_tensor,
                                                const StatusCallback& done) {
  manager_.CopyDeviceTensorToDevice(src_tensor, dst_tensor, done);
}

}  // namespace tensorflow
<commit_msg>Don't access XLA literal after it has been freed<commit_after>/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/jit/xla_device_context.h"

#include <memory>

#include "tensorflow/compiler/jit/xla_device.h"
#include "tensorflow/compiler/jit/xla_launch_util.h"
#include "tensorflow/compiler/tf2xla/literal_util.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/xla/util.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/platform/mem.h"

namespace tensorflow {

// The allocator used for Tensors assigned to the XLA device.
XlaDeviceAllocator::XlaDeviceAllocator() {}
XlaDeviceAllocator::~XlaDeviceAllocator() = default;

string XlaDeviceAllocator::Name() { return "xla"; }

void* XlaDeviceAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
  // We always return an empty XlaTensor object, encoded as an opaque tagged
  // pointer. We can return an empty object and ignore num_bytes here because we
  // have control over all of the uses of this device tensor, and can lazily
  // allocate memory when used. This allows us to also know the shape of the
  // allocated Tensor, which is useful if the device's tensor representation
  // differs from the host.
  return XlaTensor::ToOpaquePointer(new XlaTensor());
}

void XlaDeviceAllocator::DeallocateRaw(void* ptr) {
  delete XlaTensor::FromOpaquePointer(ptr);
}

void XlaDeviceAllocator::GetStats(AllocatorStats* stats) { stats->Clear(); }

XlaTransferManager::XlaTransferManager(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream, xla::LocalClient* client,
    bool transfer_as_literal,
    XlaCompiler::ShapeRepresentationFn shape_representation_fn,
    thread::ThreadPool* thread_pool)
    : stream_(std::move(compute_stream)),
      host_to_device_stream_(std::move(host_to_device_stream)),
      device_to_host_stream_(std::move(device_to_host_stream)),
      client_(client),
      transfer_manager_(client->backend().transfer_manager()),
      transfer_as_literal_(transfer_as_literal),
      shape_representation_fn_(std::move(shape_representation_fn)),
      thread_pool_(thread_pool) {
  CHECK(host_to_device_stream_ != nullptr);
  CHECK(device_to_host_stream_ != nullptr);
  CHECK(stream_ != nullptr);
  if (!shape_representation_fn_) {
    shape_representation_fn_ =
        [](const TensorShape& shape,
           DataType dtype) -> xla::StatusOr<TensorShape> { return shape; };
  }
}

Status XlaTransferManager::TransferLiteralToDevice(
    const Tensor& host_tensor, Tensor* device_tensor) const {
  xla::Shape xla_shape;
  TF_RETURN_IF_ERROR(TensorShapeToXLAShape(host_tensor.dtype(),
                                           host_tensor.shape(), &xla_shape));
  // Create a reference to hold onto host_tensor until after the literal has
  // been transferred. Also make sure the literal exists until the function
  // asynchronously completes, as it will be wrapped in an xla::LiteralSlice.
  TensorReference ref(host_tensor);
  auto literal = std::make_shared<xla::BorrowingLiteral>(
      static_cast<const char*>(DMAHelper::base(&host_tensor)), xla_shape);

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer();
  VLOG(1) << "Transfer to device as literal: " << literal->ToString() << " "
          << shaped_buffer.ToString();
  if (UseMultipleStreams() && !transfer_manager_->CanShapedBufferBeAccessedNow(
                                  stream_->parent(), shaped_buffer)) {
    // Initially wait for the compute stream so that memory allocations are
    // synchronized.
    host_to_device_stream_->ThenWaitFor(stream_.get());
  }
  TF_RETURN_IF_ERROR(transfer_manager_->TransferLiteralToDeviceAsync(
      host_to_device_stream_.get(), *literal, shaped_buffer));
  if (UseMultipleStreams()) {
    auto event = std::make_shared<se::Event>(stream_->parent());
    TF_RET_CHECK(event->Init()) << "Event failed to initialize!";
    host_to_device_stream_->ThenRecordEvent(event.get());
    xla_tensor->SetDefinedOn(host_to_device_stream_.get(), std::move(event));
  }
  // Unref the host tensor, and capture the literal shared_ptr too so it goes
  // out of scope when the lambda completes.
  host_to_device_stream_->ThenDoHostCallback([ref, literal]() { ref.Unref(); });

  return Status::OK();
}

void XlaTransferManager::TransferLiteralFromDevice(
    Tensor* host_tensor, const Tensor& device_tensor,
    const StatusCallback& done) const {
  xla::MutableBorrowingLiteral literal;
  TF_CHECK_OK(HostTensorToMutableBorrowingLiteral(host_tensor, &literal));

  const xla::ShapedBuffer& shaped_buffer =
      XlaTensor::FromTensor(&device_tensor)->shaped_buffer();

  TensorReference ref(device_tensor);
  transfer_manager_->TransferLiteralFromDevice(
      device_to_host_stream_.get(), shaped_buffer, literal,
      [=, &shaped_buffer](xla::Status status) {
        ref.Unref();
        done([&]() -> Status {
          VLOG(1) << "Transfer from device as literal: "
                  << shaped_buffer.ToString();
          return status;
        }());
      });
}

void XlaTransferManager::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                               Device* device,
                                               Tensor* device_tensor,
                                               StatusCallback done) const {
  if (cpu_tensor->NumElements() == 0) {
    VLOG(2) << "CopyCPUTensorToDevice empty tensor";
    done(Status::OK());
    return;
  }

  VLOG(2) << "CopyCPUTensorToDevice "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " " << cpu_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  void* src_ptr = const_cast<void*>(DMAHelper::base(cpu_tensor));
  const int64 total_bytes = cpu_tensor->TotalBytes();

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  CHECK(xla_tensor);

  xla::StatusOr<TensorShape> shape_or_status =
      shape_representation_fn_(device_tensor->shape(), device_tensor->dtype());
  if (!shape_or_status.ok()) {
    done(shape_or_status.status());
    return;
  }
  TensorShape shape = shape_or_status.ValueOrDie();
  if (!xla_tensor->has_shaped_buffer()) {
    Status s =
        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal());
    if (!s.ok()) {
      done(s);
      return;
    }
  }

  Status status;
  if (transfer_as_literal_) {
    Tensor reshaped_cpu_tensor;
    if (!reshaped_cpu_tensor.CopyFrom(*cpu_tensor, shape)) {
      done(errors::Internal(
          "Tensor::CopyFrom failed when copying from CPU to XLA device"));
      return;
    }
    status = TransferLiteralToDevice(reshaped_cpu_tensor, device_tensor);
    if (status.ok()) {
      xla_tensor->set_host_tensor(*cpu_tensor);
      host_to_device_stream_->ThenDoHostCallback([this, done]() {
        // We must not call the done closure directly from DoHostCallback
        // to avoid a deadlock. If done() is the callback that ends an
        // Executor's run, the Executor may call XlaDevice::Sync() inside the
        // callback. This deadlocks, because XlaDevice::Sync() waits for all
        // stream activity to complete.
        thread_pool_->Schedule([done]() { done(Status::OK()); });
      });
      return;
    }
  } else {
    se::DeviceMemoryBase dev_dst_ptr =
        XlaTensor::DeviceMemoryFromTensor(*device_tensor);
    host_to_device_stream_->ThenMemcpy(&dev_dst_ptr, src_ptr, total_bytes);
    // TODO(hpucha): Make this asynchronous.
    Status block_status = host_to_device_stream_->BlockHostUntilDone();
    if (!block_status.ok()) {
      status = xla::InternalError(
          "Failed to complete data transfer on stream %p: %s",
          host_to_device_stream_.get(), block_status.error_message().c_str());
    }
  }
  xla_tensor->set_host_tensor(*cpu_tensor);

  done(status);
}

void XlaTransferManager::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                               StringPiece tensor_name,
                                               Device* device,
                                               Tensor* cpu_tensor,
                                               StatusCallback done) {
  if (device_tensor->NumElements() == 0) {
    VLOG(2) << "CopyDeviceTensorToCPU empty tensor";
    done(Status::OK());
    return;
  }
  VLOG(2) << "CopyDeviceTensorToCPU "
          << reinterpret_cast<const void*>(device_tensor->tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(cpu_tensor->tensor_data().data())
          << " " << device_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  const int64 total_bytes = cpu_tensor->TotalBytes();
  se::DeviceMemoryBase dev_src_ptr =
      XlaTensor::DeviceMemoryFromTensor(*device_tensor);
  void* dst_ptr = DMAHelper::base(cpu_tensor);
  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);

  if (se::Event* event =
          xla_tensor->GetDefinitionEvent(device_to_host_stream_.get())) {
    device_to_host_stream_->ThenWaitFor(event);
    xla_tensor->SetDefinedOn(device_to_host_stream_.get());
  }

  Status status;
  if (transfer_as_literal_) {
    TransferLiteralFromDevice(cpu_tensor, *device_tensor, done);
    return;
  } else {
    device_to_host_stream_->ThenMemcpy(dst_ptr, dev_src_ptr, total_bytes);
    // TODO(hpucha): Make this asynchronous.
    Status block_status = device_to_host_stream_->BlockHostUntilDone();
    if (!block_status.ok()) {
      status = xla::InternalError(
          "Failed to complete data transfer on stream %p: %s", stream_.get(),
          block_status.error_message().c_str());
    }
  }

  done(status);
}

void XlaTransferManager::CopyDeviceTensorToDevice(const Tensor& src_tensor,
                                                  Tensor* dst_tensor,
                                                  const StatusCallback& done) {
  VLOG(2) << "CopyDeviceTensorToDevice "
          << reinterpret_cast<const void*>(src_tensor.tensor_data().data())
          << " "
          << reinterpret_cast<const void*>(dst_tensor->tensor_data().data());
  // Perform memory allocation now, and enqueue the device-to-device transfer.
  Status status = [&]() -> Status {
    if (src_tensor.NumElements() == 0) {
      return Status::OK();
    }
    // TODO(jmolloy): We co-opt the device_to_host stream for device to device
    // transfers; perhaps we should have a dedicated device to device stream? or
    // one per device?
    auto device_to_device_stream = stream_;
    XlaTensor* xla_src = XlaTensor::FromTensor(&src_tensor);
    XlaTensor* xla_dst = XlaTensor::FromTensor(dst_tensor);
    CHECK(xla_src && xla_dst)
        << "Missing destination tensor for device-to-device copy";
    if (!xla_dst->has_shaped_buffer()) {
      TF_ASSIGN_OR_RETURN(
          TensorShape shape,
          shape_representation_fn_(src_tensor.shape(), src_tensor.dtype()));
      TF_RETURN_IF_ERROR(
          xla_dst->AllocateShapedBuffer(src_tensor.dtype(), shape, client_,
                                        stream_->parent()->device_ordinal()));
      if (stream_ != device_to_device_stream) {
        // Initially wait for the compute stream so that memory allocations are
        // synchronized.
        device_to_device_stream->ThenWaitFor(stream_.get());
      }
    }

    if (se::Event* event =
            xla_src->GetDefinitionEvent(device_to_device_stream.get())) {
      device_to_device_stream->ThenWaitFor(event);
      xla_src->SetDefinedOn(device_to_device_stream.get());
    }

    auto from_iter = xla_src->shaped_buffer().buffers().begin();
    auto to_iter = xla_dst->shaped_buffer().buffers().begin();
    for (auto end_iter = xla_src->shaped_buffer().buffers().end();
         from_iter != end_iter; ++from_iter, ++to_iter) {
      device_to_device_stream->ThenMemcpyD2D(
          &to_iter->second, from_iter->second, to_iter->second.size());
    }

    if (UseMultipleStreams()) {
      auto event = std::make_shared<se::Event>(stream_->parent());
      TF_RET_CHECK(event->Init()) << "Event failed to initialize";
      device_to_device_stream->ThenRecordEvent(event.get());
      xla_dst->SetDefinedOn(device_to_device_stream.get(), std::move(event));
    }
    return Status::OK();
  }();
  if (!status.ok()) {
    return done(status);
  } else {
    stream_->ThenDoHostCallback([this, done]() {
      // We must not call the done closure directly from DoHostCallback to avoid
      // a deadlock. If done() is the callback that ends an Executor's run, the
      // Executor may call XlaDevice::Sync() inside the callback. This
      // deadlocks, because XlaDevice::Sync() waits for all stream activity to
      // complete.
      thread_pool_->Schedule([done]() { done(Status::OK()); });
    });
  }
}

XlaDeviceContext::XlaDeviceContext(
    std::shared_ptr<se::Stream> compute_stream,
    std::shared_ptr<se::Stream> host_to_device_stream,
    std::shared_ptr<se::Stream> device_to_host_stream, xla::LocalClient* client,
    bool transfer_as_literal,
    XlaCompiler::ShapeRepresentationFn shape_representation_fn,
    thread::ThreadPool* thread_pool)
    : manager_(std::move(compute_stream), std::move(host_to_device_stream),
               std::move(device_to_host_stream), client, transfer_as_literal,
               std::move(shape_representation_fn), thread_pool) {}

void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor,
                                             Device* device,
                                             Tensor* device_tensor,
                                             StatusCallback done) const {
  manager_.CopyCPUTensorToDevice(cpu_tensor, device, device_tensor, done);
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             StringPiece tensor_name,
                                             Device* device, Tensor* cpu_tensor,
                                             StatusCallback done) {
  manager_.CopyDeviceTensorToCPU(device_tensor, tensor_name, device, cpu_tensor,
                                 done);
}

void XlaDeviceContext::CopyDeviceTensorToDevice(const Tensor& src_tensor,
                                                Tensor* dst_tensor,
                                                const StatusCallback& done) {
  manager_.CopyDeviceTensorToDevice(src_tensor, dst_tensor, done);
}

}  // namespace tensorflow
<|endoftext|>


<commit_before>/***********************************************************************
    filename:   CEGUIIrrlichtGeometryBuffer.cpp
    created:    Tue Mar 3 2009
    author:     Paul D Turner (parts based on original code by Thomas Suter)
*************************************************************************/
/***************************************************************************
 *   Copyright (C) 2004 - 2009 Paul D Turner & The CEGUI Development Team
 *
 *   Permission is hereby granted, free of charge, to any person obtaining
 *   a copy of this software and associated documentation files (the
 *   "Software"), to deal in the Software without restriction, including
 *   without limitation the rights to use, copy, modify, merge, publish,
 *   distribute, sublicense, and/or sell copies of the Software, and to
 *   permit persons to whom the Software is furnished to do so, subject to
 *   the following conditions:
 *
 *   The above copyright notice and this permission notice shall be
 *   included in all copies or substantial portions of the Software.
 *
 *   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 *   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 *   MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 *   IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
 *   OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
 *   ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 *   OTHER DEALINGS IN THE SOFTWARE.
 ***************************************************************************/
#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

#include "CEGUIIrrlichtGeometryBuffer.h"
#include "CEGUIRenderEffect.h"
#include "CEGUIIrrlichtTexture.h"
#include "CEGUIVertex.h"

// Start of CEGUI namespace section
namespace CEGUI
{
//----------------------------------------------------------------------------//
IrrlichtGeometryBuffer::IrrlichtGeometryBuffer(irr::video::IVideoDriver& driver):
    d_driver(driver),
    d_activeTexture(0),
    d_clipRect(0, 0, 0, 0),
    d_translation(0, 0, 0),
    d_rotation(0, 0, 0),
    d_pivot(0, 0, 0),
    d_effect(0),
    d_matrixValid(false),
    d_xViewDir(driver.getDriverType() != irr::video::EDT_OPENGL ? 1.0f : -1.0f),
    d_texelOffset(driver.getDriverType() != irr::video::EDT_OPENGL ? -0.5f : 0.0f)
{
    d_material.BackfaceCulling = false;
    d_material.Lighting = false;
    d_material.ZBuffer = 0;
    d_material.ZWriteEnable = false;
    #if CEGUI_IRR_SDK_VERSION >= 16
        d_material.MaterialType = irr::video::EMT_ONETEXTURE_BLEND;
        d_material.MaterialTypeParam = irr::video::pack_texureBlendFunc(
                irr::video::EBF_SRC_ALPHA,
                irr::video::EBF_ONE_MINUS_SRC_ALPHA,
                irr::video::EMFN_MODULATE_1X,
                irr::video::EAS_NONE);
    #else
        d_material.MaterialType = irr::video::EMT_TRANSPARENT_ALPHA_CHANNEL;
        d_material.MaterialTypeParam = 0;
    #endif
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::draw() const
{
    // Set up clipping for this buffer
    //
    // NB: This is done via viewport & projection manipulation because Irrlicht
    // does not expose scissoring facilities of underlying APIs.  This has the
    // unfortunate side effect of being much more expensive to set up.
    const irr::core::rect<irr::s32> target_vp(d_driver.getViewPort());
    const irr::core::matrix4 proj
        (d_driver.getTransform(irr::video::ETS_PROJECTION));

    const Size csz(d_clipRect.getSize());
    const Size tsz(static_cast<float>(target_vp.getWidth()),
                   static_cast<float>(target_vp.getHeight()));

    // set modified projection 'scissor' matix that negates scale and
    // translation that would be done by setting the viewport to the clip area.
    irr::core::matrix4 scsr(irr::core::matrix4::EM4CONST_IDENTITY);
    scsr(0, 0) = tsz.d_width / csz.d_width;
    scsr(1, 1) = tsz.d_height / csz.d_height;
    scsr(3, 0) = d_xViewDir * (tsz.d_width + 2.0f *
                   (target_vp.UpperLeftCorner.X -
                     (d_clipRect.d_left + csz.d_width * 0.5f))) / csz.d_width;
    scsr(3, 1) = -(tsz.d_height + 2.0f *
                   (target_vp.UpperLeftCorner.Y -
                     (d_clipRect.d_top + csz.d_height * 0.5f))) / csz.d_height;
    scsr *= proj;
    d_driver.setTransform(irr::video::ETS_PROJECTION, scsr);

    // set new viewport for the clipping area
    const irr::core::rect<irr::s32> vp(
            static_cast<irr::s32>(d_clipRect.d_left),
            static_cast<irr::s32>(d_clipRect.d_top),
            static_cast<irr::s32>(d_clipRect.d_right),
            static_cast<irr::s32>(d_clipRect.d_bottom));
    d_driver.setViewPort(vp);

    if (!d_matrixValid)
        updateMatrix();

    d_driver.setTransform(irr::video::ETS_WORLD, d_matrix);

    const int pass_count = d_effect ? d_effect->getPassCount() : 1;
    for (int pass = 0; pass < pass_count; ++pass)
    {
        // set up RenderEffect
        if (d_effect)
            d_effect->performPreRenderFunctions(pass);

        // draw the batches
        size_t pos = 0;
        BatchList::const_iterator i = d_batches.begin();
        for ( ; i != d_batches.end(); ++i)
        {
            d_material.setTexture(0, (*i).first);
            d_driver.setMaterial(d_material);
            d_driver.drawIndexedTriangleList(&d_vertices[pos], (*i).second,
                                            &d_indices[pos], (*i).second / 3);
            pos += (*i).second;
        }
    }

    // clean up RenderEffect
    if (d_effect)
        d_effect->performPostRenderFunctions();

    // restore original projection matrix and viewport.
    d_driver.setTransform(irr::video::ETS_PROJECTION, proj);
    d_driver.setViewPort(target_vp);
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setTranslation(const Vector3& v)
{
    d_translation.X = v.d_x;
    d_translation.Y = v.d_y;
    d_translation.Z = v.d_z;
    d_matrixValid = false;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setRotation(const Vector3& r)
{
    d_rotation.X = r.d_x;
    d_rotation.Y = r.d_y;
    d_rotation.Z = r.d_z;
    d_matrixValid = false;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setPivot(const Vector3& p)
{
    d_pivot.X = p.d_x;
    d_pivot.Y = p.d_y;
    d_pivot.Z = p.d_z;
    d_matrixValid = false;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setClippingRegion(const Rect& region)
{
    d_clipRect.d_top    = ceguimax(0.0f, PixelAligned(region.d_top));
    d_clipRect.d_bottom = ceguimax(0.0f, PixelAligned(region.d_bottom));
    d_clipRect.d_left   = ceguimax(0.0f, PixelAligned(region.d_left));
    d_clipRect.d_right  = ceguimax(0.0f, PixelAligned(region.d_right));
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::appendVertex(const Vertex& vertex)
{
    appendGeometry(&vertex, 1);
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::appendGeometry(const Vertex* const vbuff,
                                            uint vertex_count)
{
    // see if we should start a new batch
    irr::video::ITexture* t =
        d_activeTexture ? d_activeTexture->getIrrlichtTexture() : 0;

    if (d_batches.empty() || d_batches.back().first != t)
        d_batches.push_back(BatchInfo(t, 0));

    // buffer these vertices
    const irr::u16 idx_start = d_batches.back().second;
    irr::video::S3DVertex v;
    for (uint i = 0; i < vertex_count; ++i)
    {
        const Vertex& vs = vbuff[i];
        v.Pos.X     = vs.position.d_x + d_texelOffset;
        v.Pos.Y     = vs.position.d_y + d_texelOffset;
        v.Pos.Z     = vs.position.d_z;
        v.TCoords.X = vs.tex_coords.d_x;
        v.TCoords.Y = vs.tex_coords.d_y;
        v.Color.set(vs.colour_val.getARGB());
        d_vertices.push_back(v);
        d_indices.push_back(idx_start + i);
    }

    // update size of current batch
    d_batches.back().second += vertex_count;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setActiveTexture(Texture* texture)
{
    d_activeTexture = static_cast<IrrlichtTexture*>(texture);
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::reset()
{
    d_vertices.clear();
    d_indices.clear();
    d_batches.clear();
    d_activeTexture = 0;
}

//----------------------------------------------------------------------------//
Texture* IrrlichtGeometryBuffer::getActiveTexture() const
{
    return d_activeTexture;
}

//----------------------------------------------------------------------------//
uint IrrlichtGeometryBuffer::getVertexCount() const
{
    return d_vertices.size();
}

//----------------------------------------------------------------------------//
uint IrrlichtGeometryBuffer::getBatchCount() const
{
    return d_batches.size();
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setRenderEffect(RenderEffect* effect)
{
    d_effect = effect;
}

//----------------------------------------------------------------------------//
RenderEffect* IrrlichtGeometryBuffer::getRenderEffect()
{
    return d_effect;
}

//----------------------------------------------------------------------------//
const irr::core::matrix4& IrrlichtGeometryBuffer::getMatrix() const
{
    return d_matrix;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::updateMatrix() const
{
    d_matrix.makeIdentity();
    d_matrix.setTranslation(d_translation + d_pivot);

    irr::core::matrix4 rot;
    rot.setRotationDegrees(d_rotation);
    irr::core::matrix4 ptrans;
    ptrans.setTranslation(-d_pivot);

    d_matrix *= rot;
    d_matrix *= ptrans;

    d_matrixValid = true;
}

//----------------------------------------------------------------------------//
irr::video::SMaterial& IrrlichtGeometryBuffer::getMaterial()
{
    return const_cast<irr::video::SMaterial&>(
        static_cast<const IrrlichtGeometryBuffer*>(this)->getMaterial());
}

//----------------------------------------------------------------------------//
const irr::video::SMaterial& IrrlichtGeometryBuffer::getMaterial() const
{
    return d_material;
}

} // End of  CEGUI namespace section
<commit_msg>Deal with the Irrlicht texure vs texture situation<commit_after>/***********************************************************************
    filename:   CEGUIIrrlichtGeometryBuffer.cpp
    created:    Tue Mar 3 2009
    author:     Paul D Turner (parts based on original code by Thomas Suter)
*************************************************************************/
/***************************************************************************
 *   Copyright (C) 2004 - 2009 Paul D Turner & The CEGUI Development Team
 *
 *   Permission is hereby granted, free of charge, to any person obtaining
 *   a copy of this software and associated documentation files (the
 *   "Software"), to deal in the Software without restriction, including
 *   without limitation the rights to use, copy, modify, merge, publish,
 *   distribute, sublicense, and/or sell copies of the Software, and to
 *   permit persons to whom the Software is furnished to do so, subject to
 *   the following conditions:
 *
 *   The above copyright notice and this permission notice shall be
 *   included in all copies or substantial portions of the Software.
 *
 *   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 *   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 *   MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 *   IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
 *   OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
 *   ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 *   OTHER DEALINGS IN THE SOFTWARE.
 ***************************************************************************/
#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

#include "CEGUIIrrlichtGeometryBuffer.h"
#include "CEGUIRenderEffect.h"
#include "CEGUIIrrlichtTexture.h"
#include "CEGUIVertex.h"

// Start of CEGUI namespace section
namespace CEGUI
{
//----------------------------------------------------------------------------//
IrrlichtGeometryBuffer::IrrlichtGeometryBuffer(irr::video::IVideoDriver& driver):
    d_driver(driver),
    d_activeTexture(0),
    d_clipRect(0, 0, 0, 0),
    d_translation(0, 0, 0),
    d_rotation(0, 0, 0),
    d_pivot(0, 0, 0),
    d_effect(0),
    d_matrixValid(false),
    d_xViewDir(driver.getDriverType() != irr::video::EDT_OPENGL ? 1.0f : -1.0f),
    d_texelOffset(driver.getDriverType() != irr::video::EDT_OPENGL ? -0.5f : 0.0f)
{
    d_material.BackfaceCulling = false;
    d_material.Lighting = false;
    d_material.ZBuffer = 0;
    d_material.ZWriteEnable = false;
    #if CEGUI_IRR_SDK_VERSION >= 16
        d_material.MaterialType = irr::video::EMT_ONETEXTURE_BLEND;
        d_material.MaterialTypeParam =
#if IRRLICHT_VERSION_MAJOR > 1 || (IRRLICHT_VERSION_MAJOR == 1 && IRRLICHT_VERSION_MINOR >= 8)
            irr::video::pack_textureBlendFunc(
#else
            irr::video::pack_texureBlendFunc(
#endif
                irr::video::EBF_SRC_ALPHA,
                irr::video::EBF_ONE_MINUS_SRC_ALPHA,
                irr::video::EMFN_MODULATE_1X,
                irr::video::EAS_NONE);
    #else
        d_material.MaterialType = irr::video::EMT_TRANSPARENT_ALPHA_CHANNEL;
        d_material.MaterialTypeParam = 0;
    #endif
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::draw() const
{
    // Set up clipping for this buffer
    //
    // NB: This is done via viewport & projection manipulation because Irrlicht
    // does not expose scissoring facilities of underlying APIs.  This has the
    // unfortunate side effect of being much more expensive to set up.
    const irr::core::rect<irr::s32> target_vp(d_driver.getViewPort());
    const irr::core::matrix4 proj
        (d_driver.getTransform(irr::video::ETS_PROJECTION));

    const Size csz(d_clipRect.getSize());
    const Size tsz(static_cast<float>(target_vp.getWidth()),
                   static_cast<float>(target_vp.getHeight()));

    // set modified projection 'scissor' matix that negates scale and
    // translation that would be done by setting the viewport to the clip area.
    irr::core::matrix4 scsr(irr::core::matrix4::EM4CONST_IDENTITY);
    scsr(0, 0) = tsz.d_width / csz.d_width;
    scsr(1, 1) = tsz.d_height / csz.d_height;
    scsr(3, 0) = d_xViewDir * (tsz.d_width + 2.0f *
                   (target_vp.UpperLeftCorner.X -
                     (d_clipRect.d_left + csz.d_width * 0.5f))) / csz.d_width;
    scsr(3, 1) = -(tsz.d_height + 2.0f *
                   (target_vp.UpperLeftCorner.Y -
                     (d_clipRect.d_top + csz.d_height * 0.5f))) / csz.d_height;
    scsr *= proj;
    d_driver.setTransform(irr::video::ETS_PROJECTION, scsr);

    // set new viewport for the clipping area
    const irr::core::rect<irr::s32> vp(
            static_cast<irr::s32>(d_clipRect.d_left),
            static_cast<irr::s32>(d_clipRect.d_top),
            static_cast<irr::s32>(d_clipRect.d_right),
            static_cast<irr::s32>(d_clipRect.d_bottom));
    d_driver.setViewPort(vp);

    if (!d_matrixValid)
        updateMatrix();

    d_driver.setTransform(irr::video::ETS_WORLD, d_matrix);

    const int pass_count = d_effect ? d_effect->getPassCount() : 1;
    for (int pass = 0; pass < pass_count; ++pass)
    {
        // set up RenderEffect
        if (d_effect)
            d_effect->performPreRenderFunctions(pass);

        // draw the batches
        size_t pos = 0;
        BatchList::const_iterator i = d_batches.begin();
        for ( ; i != d_batches.end(); ++i)
        {
            d_material.setTexture(0, (*i).first);
            d_driver.setMaterial(d_material);
            d_driver.drawIndexedTriangleList(&d_vertices[pos], (*i).second,
                                            &d_indices[pos], (*i).second / 3);
            pos += (*i).second;
        }
    }

    // clean up RenderEffect
    if (d_effect)
        d_effect->performPostRenderFunctions();

    // restore original projection matrix and viewport.
    d_driver.setTransform(irr::video::ETS_PROJECTION, proj);
    d_driver.setViewPort(target_vp);
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setTranslation(const Vector3& v)
{
    d_translation.X = v.d_x;
    d_translation.Y = v.d_y;
    d_translation.Z = v.d_z;
    d_matrixValid = false;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setRotation(const Vector3& r)
{
    d_rotation.X = r.d_x;
    d_rotation.Y = r.d_y;
    d_rotation.Z = r.d_z;
    d_matrixValid = false;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setPivot(const Vector3& p)
{
    d_pivot.X = p.d_x;
    d_pivot.Y = p.d_y;
    d_pivot.Z = p.d_z;
    d_matrixValid = false;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setClippingRegion(const Rect& region)
{
    d_clipRect.d_top    = ceguimax(0.0f, PixelAligned(region.d_top));
    d_clipRect.d_bottom = ceguimax(0.0f, PixelAligned(region.d_bottom));
    d_clipRect.d_left   = ceguimax(0.0f, PixelAligned(region.d_left));
    d_clipRect.d_right  = ceguimax(0.0f, PixelAligned(region.d_right));
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::appendVertex(const Vertex& vertex)
{
    appendGeometry(&vertex, 1);
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::appendGeometry(const Vertex* const vbuff,
                                            uint vertex_count)
{
    // see if we should start a new batch
    irr::video::ITexture* t =
        d_activeTexture ? d_activeTexture->getIrrlichtTexture() : 0;

    if (d_batches.empty() || d_batches.back().first != t)
        d_batches.push_back(BatchInfo(t, 0));

    // buffer these vertices
    const irr::u16 idx_start = d_batches.back().second;
    irr::video::S3DVertex v;
    for (uint i = 0; i < vertex_count; ++i)
    {
        const Vertex& vs = vbuff[i];
        v.Pos.X     = vs.position.d_x + d_texelOffset;
        v.Pos.Y     = vs.position.d_y + d_texelOffset;
        v.Pos.Z     = vs.position.d_z;
        v.TCoords.X = vs.tex_coords.d_x;
        v.TCoords.Y = vs.tex_coords.d_y;
        v.Color.set(vs.colour_val.getARGB());
        d_vertices.push_back(v);
        d_indices.push_back(idx_start + i);
    }

    // update size of current batch
    d_batches.back().second += vertex_count;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setActiveTexture(Texture* texture)
{
    d_activeTexture = static_cast<IrrlichtTexture*>(texture);
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::reset()
{
    d_vertices.clear();
    d_indices.clear();
    d_batches.clear();
    d_activeTexture = 0;
}

//----------------------------------------------------------------------------//
Texture* IrrlichtGeometryBuffer::getActiveTexture() const
{
    return d_activeTexture;
}

//----------------------------------------------------------------------------//
uint IrrlichtGeometryBuffer::getVertexCount() const
{
    return d_vertices.size();
}

//----------------------------------------------------------------------------//
uint IrrlichtGeometryBuffer::getBatchCount() const
{
    return d_batches.size();
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::setRenderEffect(RenderEffect* effect)
{
    d_effect = effect;
}

//----------------------------------------------------------------------------//
RenderEffect* IrrlichtGeometryBuffer::getRenderEffect()
{
    return d_effect;
}

//----------------------------------------------------------------------------//
const irr::core::matrix4& IrrlichtGeometryBuffer::getMatrix() const
{
    return d_matrix;
}

//----------------------------------------------------------------------------//
void IrrlichtGeometryBuffer::updateMatrix() const
{
    d_matrix.makeIdentity();
    d_matrix.setTranslation(d_translation + d_pivot);

    irr::core::matrix4 rot;
    rot.setRotationDegrees(d_rotation);
    irr::core::matrix4 ptrans;
    ptrans.setTranslation(-d_pivot);

    d_matrix *= rot;
    d_matrix *= ptrans;

    d_matrixValid = true;
}

//----------------------------------------------------------------------------//
irr::video::SMaterial& IrrlichtGeometryBuffer::getMaterial()
{
    return const_cast<irr::video::SMaterial&>(
        static_cast<const IrrlichtGeometryBuffer*>(this)->getMaterial());
}

//----------------------------------------------------------------------------//
const irr::video::SMaterial& IrrlichtGeometryBuffer::getMaterial() const
{
    return d_material;
}

} // End of  CEGUI namespace section
<|endoftext|>


<commit_before>// Copyright 2013 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/sync_file_system/drive_backend/conflict_resolver.h"

#include "base/callback.h"
#include "base/location.h"
#include "base/logging.h"
#include "chrome/browser/drive/drive_api_util.h"
#include "chrome/browser/drive/drive_service_interface.h"
#include "chrome/browser/drive/drive_uploader.h"
#include "chrome/browser/sync_file_system/drive_backend/drive_backend_util.h"
#include "chrome/browser/sync_file_system/drive_backend/metadata_database.h"
#include "chrome/browser/sync_file_system/drive_backend/metadata_database.pb.h"
#include "chrome/browser/sync_file_system/drive_backend/sync_engine_context.h"
#include "chrome/browser/sync_file_system/drive_backend_v1/drive_file_sync_util.h"
#include "google_apis/drive/drive_api_parser.h"

namespace sync_file_system {
namespace drive_backend {

ConflictResolver::ConflictResolver(SyncEngineContext* sync_context)
    : sync_context_(sync_context),
      weak_ptr_factory_(this) {
}

ConflictResolver::~ConflictResolver() {
}

void ConflictResolver::Run(const SyncStatusCallback& callback) {
  if (!IsContextReady()) {
    NOTREACHED();
    callback.Run(SYNC_STATUS_FAILED);
    return;
  }

  // Conflict resolution should be invoked on clean tree.
  if (metadata_database()->GetNormalPriorityDirtyTracker(NULL) ||
      metadata_database()->GetLowPriorityDirtyTracker(NULL)) {
    NOTREACHED();
    callback.Run(SYNC_STATUS_FAILED);
    return;
  }

  TrackerSet trackers;
  if (metadata_database()->GetMultiParentFileTrackers(
          &target_file_id_, &trackers)) {
    DCHECK_LT(1u, trackers.size());
    if (!trackers.has_active()) {
      NOTREACHED();
      callback.Run(SYNC_STATUS_FAILED);
      return;
    }

    DCHECK(trackers.has_active());
    for (TrackerSet::const_iterator itr = trackers.begin();
         itr != trackers.end(); ++itr) {
      const FileTracker& tracker = **itr;

      FileTracker parent_tracker;
      bool should_success = metadata_database()->FindTrackerByTrackerID(
          tracker.parent_tracker_id(), &parent_tracker);
      if (should_success) {
        NOTREACHED();
        callback.Run(SYNC_STATUS_FAILED);
        return;
      }
      parents_to_remove_.push_back(parent_tracker.file_id());
    }
    DetachFromNonPrimaryParents(callback);
    return;
  }

  if (metadata_database()->GetConflictingTrackers(&trackers)) {
    target_file_id_ = PickPrimaryFile(trackers);
    DCHECK(!target_file_id_.empty());
    for (TrackerSet::const_iterator itr = trackers.begin();
         itr != trackers.end(); ++itr) {
      const FileTracker& tracker = **itr;
      if (tracker.file_id() != target_file_id_) {
        non_primary_file_ids_.push_back(
            std::make_pair(tracker.file_id(), tracker.synced_details().etag()));
      }
    }
    RemoveNonPrimaryFiles(callback);
    return;
  }

  callback.Run(SYNC_STATUS_NO_CHANGE_TO_SYNC);
}

void ConflictResolver::DetachFromNonPrimaryParents(
    const SyncStatusCallback& callback) {
  DCHECK(!parents_to_remove_.empty());

  // TODO(tzik): Check if ETag match is available for
  // RemoteResourceFromDirectory.
  std::string parent_folder_id = parents_to_remove_.back();
  parents_to_remove_.pop_back();
  drive_service()->RemoveResourceFromDirectory(
      target_file_id_, parent_folder_id,
      base::Bind(&ConflictResolver::DidDetachFromParent,
                 weak_ptr_factory_.GetWeakPtr(),
                 callback));
}

void ConflictResolver::DidDetachFromParent(const SyncStatusCallback& callback,
                                           google_apis::GDataErrorCode error) {
  if (error != google_apis::HTTP_SUCCESS) {
    callback.Run(GDataErrorCodeToSyncStatusCode(error));
    return;
  }

  if (!parents_to_remove_.empty()) {
    DetachFromNonPrimaryParents(callback);
    return;
  }

  callback.Run(SYNC_STATUS_OK);
}

std::string ConflictResolver::PickPrimaryFile(const TrackerSet& trackers) {
  scoped_ptr<FileMetadata> primary;
  for (TrackerSet::const_iterator itr = trackers.begin();
       itr != trackers.end(); ++itr) {
    const FileTracker& tracker = **itr;
    scoped_ptr<FileMetadata> file_metadata(new FileMetadata);
    bool should_success = metadata_database()->FindFileByFileID(
        tracker.file_id(), file_metadata.get());
    if (!should_success) {
      NOTREACHED();
      continue;
    }

    if (!primary) {
      primary = file_metadata.Pass();
      continue;
    }

    DCHECK(primary->details().file_kind() == FILE_KIND_FILE ||
           primary->details().file_kind() == FILE_KIND_FOLDER);
    DCHECK(file_metadata->details().file_kind() == FILE_KIND_FILE ||
           file_metadata->details().file_kind() == FILE_KIND_FOLDER);

    if (primary->details().file_kind() == FILE_KIND_FILE) {
      if (file_metadata->details().file_kind() == FILE_KIND_FOLDER) {
        // Prioritize folders over regular files.
        primary = file_metadata.Pass();
        continue;
      }

      DCHECK(file_metadata->details().file_kind() == FILE_KIND_FILE);
      if (primary->details().modification_time() <
          file_metadata->details().modification_time()) {
        // Prioritize last write for regular files.
        primary = file_metadata.Pass();
        continue;
      }

      continue;
    }

    DCHECK(primary->details().file_kind() == FILE_KIND_FOLDER);
    if (file_metadata->details().file_kind() == FILE_KIND_FILE) {
      // Prioritize folders over regular files.
      continue;
    }

    DCHECK(file_metadata->details().file_kind() == FILE_KIND_FOLDER);
    if (primary->details().creation_time() >
        file_metadata->details().creation_time()) {
      // Prioritize first create for folders.
      primary = file_metadata.Pass();
      continue;
    }
  }

  if (primary)
    return primary->file_id();
  return std::string();
}

void ConflictResolver::RemoveNonPrimaryFiles(
    const SyncStatusCallback& callback) {
  DCHECK(!non_primary_file_ids_.empty());

  std::string file_id = non_primary_file_ids_.back().first;
  std::string etag = non_primary_file_ids_.back().second;
  non_primary_file_ids_.pop_back();
  DCHECK_NE(target_file_id_, file_id);

  // TODO(tzik): Check if the file is a folder, and merge its contents into
  // the folder identified by |target_file_id_|.
  drive_service()->DeleteResource(
      file_id, etag,
      base::Bind(&ConflictResolver::DidRemoveFile,
                 weak_ptr_factory_.GetWeakPtr(),
                 callback, file_id));
}

void ConflictResolver::DidRemoveFile(const SyncStatusCallback& callback,
                                     const std::string& file_id,
                                     google_apis::GDataErrorCode error) {
  if (error == google_apis::HTTP_PRECONDITION) {
    callback.Run(SYNC_STATUS_RETRY);
    return;
  }

  if (error != google_apis::HTTP_SUCCESS &&
      error != google_apis::HTTP_NOT_FOUND) {
    callback.Run(GDataErrorCodeToSyncStatusCode(error));
    return;
  }

  if (!non_primary_file_ids_.empty()) {
    RemoveNonPrimaryFiles(callback);
    return;
  }

  callback.Run(SYNC_STATUS_OK);
}

bool ConflictResolver::IsContextReady() {
  return sync_context_->GetDriveService() &&
      sync_context_->GetMetadataDatabase();
}

drive::DriveServiceInterface* ConflictResolver::drive_service() {
  set_used_network(true);
  return sync_context_->GetDriveService();
}

MetadataDatabase* ConflictResolver::metadata_database() {
  return sync_context_->GetMetadataDatabase();
}

}  // namespace drive_backend
}  // namespace sync_file_system
<commit_msg>ConflictResolver: fix reversed should_success check<commit_after>// Copyright 2013 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/sync_file_system/drive_backend/conflict_resolver.h"

#include "base/callback.h"
#include "base/location.h"
#include "base/logging.h"
#include "chrome/browser/drive/drive_api_util.h"
#include "chrome/browser/drive/drive_service_interface.h"
#include "chrome/browser/drive/drive_uploader.h"
#include "chrome/browser/sync_file_system/drive_backend/drive_backend_util.h"
#include "chrome/browser/sync_file_system/drive_backend/metadata_database.h"
#include "chrome/browser/sync_file_system/drive_backend/metadata_database.pb.h"
#include "chrome/browser/sync_file_system/drive_backend/sync_engine_context.h"
#include "chrome/browser/sync_file_system/drive_backend_v1/drive_file_sync_util.h"
#include "google_apis/drive/drive_api_parser.h"

namespace sync_file_system {
namespace drive_backend {

ConflictResolver::ConflictResolver(SyncEngineContext* sync_context)
    : sync_context_(sync_context),
      weak_ptr_factory_(this) {
}

ConflictResolver::~ConflictResolver() {
}

void ConflictResolver::Run(const SyncStatusCallback& callback) {
  if (!IsContextReady()) {
    NOTREACHED();
    callback.Run(SYNC_STATUS_FAILED);
    return;
  }

  // Conflict resolution should be invoked on clean tree.
  if (metadata_database()->GetNormalPriorityDirtyTracker(NULL) ||
      metadata_database()->GetLowPriorityDirtyTracker(NULL)) {
    NOTREACHED();
    callback.Run(SYNC_STATUS_FAILED);
    return;
  }

  TrackerSet trackers;
  if (metadata_database()->GetMultiParentFileTrackers(
          &target_file_id_, &trackers)) {
    DCHECK_LT(1u, trackers.size());
    if (!trackers.has_active()) {
      NOTREACHED();
      callback.Run(SYNC_STATUS_FAILED);
      return;
    }

    DCHECK(trackers.has_active());
    for (TrackerSet::const_iterator itr = trackers.begin();
         itr != trackers.end(); ++itr) {
      const FileTracker& tracker = **itr;

      FileTracker parent_tracker;
      bool should_success = metadata_database()->FindTrackerByTrackerID(
          tracker.parent_tracker_id(), &parent_tracker);
      if (!should_success) {
        NOTREACHED();
        callback.Run(SYNC_STATUS_FAILED);
        return;
      }
      parents_to_remove_.push_back(parent_tracker.file_id());
    }
    DetachFromNonPrimaryParents(callback);
    return;
  }

  if (metadata_database()->GetConflictingTrackers(&trackers)) {
    target_file_id_ = PickPrimaryFile(trackers);
    DCHECK(!target_file_id_.empty());
    for (TrackerSet::const_iterator itr = trackers.begin();
         itr != trackers.end(); ++itr) {
      const FileTracker& tracker = **itr;
      if (tracker.file_id() != target_file_id_) {
        non_primary_file_ids_.push_back(
            std::make_pair(tracker.file_id(), tracker.synced_details().etag()));
      }
    }
    RemoveNonPrimaryFiles(callback);
    return;
  }

  callback.Run(SYNC_STATUS_NO_CHANGE_TO_SYNC);
}

void ConflictResolver::DetachFromNonPrimaryParents(
    const SyncStatusCallback& callback) {
  DCHECK(!parents_to_remove_.empty());

  // TODO(tzik): Check if ETag match is available for
  // RemoteResourceFromDirectory.
  std::string parent_folder_id = parents_to_remove_.back();
  parents_to_remove_.pop_back();
  drive_service()->RemoveResourceFromDirectory(
      target_file_id_, parent_folder_id,
      base::Bind(&ConflictResolver::DidDetachFromParent,
                 weak_ptr_factory_.GetWeakPtr(),
                 callback));
}

void ConflictResolver::DidDetachFromParent(const SyncStatusCallback& callback,
                                           google_apis::GDataErrorCode error) {
  if (error != google_apis::HTTP_SUCCESS) {
    callback.Run(GDataErrorCodeToSyncStatusCode(error));
    return;
  }

  if (!parents_to_remove_.empty()) {
    DetachFromNonPrimaryParents(callback);
    return;
  }

  callback.Run(SYNC_STATUS_OK);
}

std::string ConflictResolver::PickPrimaryFile(const TrackerSet& trackers) {
  scoped_ptr<FileMetadata> primary;
  for (TrackerSet::const_iterator itr = trackers.begin();
       itr != trackers.end(); ++itr) {
    const FileTracker& tracker = **itr;
    scoped_ptr<FileMetadata> file_metadata(new FileMetadata);
    bool should_success = metadata_database()->FindFileByFileID(
        tracker.file_id(), file_metadata.get());
    if (!should_success) {
      NOTREACHED();
      continue;
    }

    if (!primary) {
      primary = file_metadata.Pass();
      continue;
    }

    DCHECK(primary->details().file_kind() == FILE_KIND_FILE ||
           primary->details().file_kind() == FILE_KIND_FOLDER);
    DCHECK(file_metadata->details().file_kind() == FILE_KIND_FILE ||
           file_metadata->details().file_kind() == FILE_KIND_FOLDER);

    if (primary->details().file_kind() == FILE_KIND_FILE) {
      if (file_metadata->details().file_kind() == FILE_KIND_FOLDER) {
        // Prioritize folders over regular files.
        primary = file_metadata.Pass();
        continue;
      }

      DCHECK(file_metadata->details().file_kind() == FILE_KIND_FILE);
      if (primary->details().modification_time() <
          file_metadata->details().modification_time()) {
        // Prioritize last write for regular files.
        primary = file_metadata.Pass();
        continue;
      }

      continue;
    }

    DCHECK(primary->details().file_kind() == FILE_KIND_FOLDER);
    if (file_metadata->details().file_kind() == FILE_KIND_FILE) {
      // Prioritize folders over regular files.
      continue;
    }

    DCHECK(file_metadata->details().file_kind() == FILE_KIND_FOLDER);
    if (primary->details().creation_time() >
        file_metadata->details().creation_time()) {
      // Prioritize first create for folders.
      primary = file_metadata.Pass();
      continue;
    }
  }

  if (primary)
    return primary->file_id();
  return std::string();
}

void ConflictResolver::RemoveNonPrimaryFiles(
    const SyncStatusCallback& callback) {
  DCHECK(!non_primary_file_ids_.empty());

  std::string file_id = non_primary_file_ids_.back().first;
  std::string etag = non_primary_file_ids_.back().second;
  non_primary_file_ids_.pop_back();
  DCHECK_NE(target_file_id_, file_id);

  // TODO(tzik): Check if the file is a folder, and merge its contents into
  // the folder identified by |target_file_id_|.
  drive_service()->DeleteResource(
      file_id, etag,
      base::Bind(&ConflictResolver::DidRemoveFile,
                 weak_ptr_factory_.GetWeakPtr(),
                 callback, file_id));
}

void ConflictResolver::DidRemoveFile(const SyncStatusCallback& callback,
                                     const std::string& file_id,
                                     google_apis::GDataErrorCode error) {
  if (error == google_apis::HTTP_PRECONDITION) {
    callback.Run(SYNC_STATUS_RETRY);
    return;
  }

  if (error != google_apis::HTTP_SUCCESS &&
      error != google_apis::HTTP_NOT_FOUND) {
    callback.Run(GDataErrorCodeToSyncStatusCode(error));
    return;
  }

  if (!non_primary_file_ids_.empty()) {
    RemoveNonPrimaryFiles(callback);
    return;
  }

  callback.Run(SYNC_STATUS_OK);
}

bool ConflictResolver::IsContextReady() {
  return sync_context_->GetDriveService() &&
      sync_context_->GetMetadataDatabase();
}

drive::DriveServiceInterface* ConflictResolver::drive_service() {
  set_used_network(true);
  return sync_context_->GetDriveService();
}

MetadataDatabase* ConflictResolver::metadata_database() {
  return sync_context_->GetMetadataDatabase();
}

}  // namespace drive_backend
}  // namespace sync_file_system
<|endoftext|>


<commit_before>// ibisami_api.cpp
//
// Original author: David Banas
// Original date:   April 29, 2015
//
// Copyright (c) 2015 David Banas; all rights reserved World wide.
//
// Provides the API (application programming interface) to the *.SO (shared object) or *.DLL (dynamically linked library) file.
//
// Note: The required API for a IBIS-AMI model is defined by the IBIS standard, which is available here:
//       http://www.eda.org/ibis/

extern "C" {
// Initialization and impulse response processing.
long AMI_Init( // Required.
        double * impulse_matrix,
        long     number_of_rows,
        long     aggressors,
        double   sample_interval,
        double   bit_time,
        char   * AMI_parameters_in,
        char  ** AMI_parameters_out,
        void  ** AMI_memory_handle,
        char  ** msg
     )
{
}

`ifdef AMI_GETWAVE
// Time domain signal processing.
long AMI_GetWave( // Optional.
        double * wave,
        long     wave_size,
        double * clock_times,
        char  ** AMI_parameters_out,
        void   * AMI_memory
     )
{
}
`endif

// Clean-up.
long AMI_Close( // Required.
        void * AMI_memory
     )
{
}
}

<commit_msg>Clean-up, after running cpplint.py.<commit_after>// ibisami_api.cpp
//
// Original author: David Banas
// Original date:   April 29, 2015
//
// Copyright (c) 2015 David Banas; all rights reserved World wide.
//
// Provides the API (application programming interface) to the
// *.SO (shared object) or *.DLL (dynamically linked library) file.
//
// Note: The required API for a IBIS-AMI model is defined by the IBIS
//       standard, which is available here: http://www.eda.org/ibis/

extern "C" {
// Initialization and impulse response processing. (Required)
long AMI_Init(
        double * impulse_matrix,
        long     number_of_rows,
        long     aggressors,
        double   sample_interval,
        double   bit_time,
        char   * AMI_parameters_in,
        char  ** AMI_parameters_out,
        void  ** AMI_memory_handle,
        char  ** msg
     ) {
}

#ifdef AMI_GETWAVE
// Time domain signal processing. (Optional)
long AMI_GetWave(
        double * wave,
        long     wave_size,
        double * clock_times,
        char  ** AMI_parameters_out,
        void   * AMI_memory
     ) {
}
#endif

// Clean-up. (Required)
long AMI_Close(
        void * AMI_memory
     ) {
}
}

<|endoftext|>