Program Listing for File tensor.h

↰ Return to documentation for file (src/tensors/tensor.h)

#pragma once

#include "common/definitions.h"
#include "common/shape.h"
#include "common/types.h"
#include "tensors/backend.h"
#include "tensors/memory_piece.h"
#ifdef CUDA_FOUND
#include "tensors/gpu/algorithm.h"
#endif

#include <algorithm>
#include <iomanip>
#include <iostream>
#include <memory>
#include <sstream>

namespace marian {

namespace io {
  struct Item;
}

class TensorBase {
  MemoryPiece::PtrType memory_;
  Shape shape_;
  Type type_{Type::float32};
  Ptr<Backend> backend_;

  ENABLE_INTRUSIVE_PTR(TensorBase)

protected:
  // Constructors are protected, use TensorBase::New(...)
  TensorBase(MemoryPiece::PtrType memory,
             Shape shape,
             Type type,
             Ptr<Backend> backend)
      : memory_(memory), shape_(shape), type_(type), backend_(backend) {}

  TensorBase(MemoryPiece::PtrType memory,
             Shape shape,
             Ptr<Backend> backend)
      : memory_(memory),
        shape_(shape),
        type_(Type::float32),
        backend_(backend) {}

  // Wraps existing memory
  template <typename T>
  TensorBase(T* rawMemory,
             size_t rawMemoryNum,
             Shape shape,
             Type type,
             Ptr<Backend> backend)
      : memory_(MemoryPiece::New((uint8_t*)rawMemory, rawMemoryNum * sizeof(T))),
        shape_(shape), type_(type), backend_(backend) {}

public:
  // Use this whenever pointing to TensorBase
  typedef IPtr<TensorBase> PtrType;

  // Use this whenever creating a pointer to TensorBase
  template <class ...Args>
  static PtrType New(Args&& ...args) {
    return PtrType(new TensorBase(std::forward<Args>(args)...));
  }

  virtual ~TensorBase() {}

  virtual void reset(MemoryPiece::PtrType memory) { memory_ = memory; }

  virtual MemoryPiece::PtrType memory() { return memory_; }

  virtual Type type() { return type_; }

  virtual Shape& shape() { return shape_; }

  virtual float* data() { return memory_->data<float>(); }

  template <typename T>
  T* data() {
    return memory_->data<T>();
  }

  virtual size_t size() { return shape_.elements(); }

  // this version of scalar will abort if numeric types do not match
  template <typename T>
  T scalar() {
    ABORT_IF(size() != 1, "Tensor is not a scalar");
    return get<T>(0);
  }

  // this non-template version converts all numeric types to float
  virtual float scalar() {
    DISPATCH_BY_TYPE0(type_, (float)scalar);
  }

  Ptr<Backend> getBackend() { return backend_; }
  DeviceId getDeviceId() { return backend_->getDeviceId(); }

  Tensor subtensor(size_t offset, size_t size) {
    auto mem = MemoryPiece::New(memory_->data() + sizeOf(type_) * offset, sizeOf(type_) * size);
    return TensorBase::New(mem, Shape{1, (int)size}, type(), backend_);
  }

  // @TODO: review if we can eliminate GPU-specific code here,
  // potentially by moving this to non-class members.
  template <typename T>
  T get(size_t i) {
    if(!matchType<T>(type_)) {
      DISPATCH_BY_TYPE1(type_, (T)get, i);
    } else {
      T temp = 0;
      if(backend_->getDeviceId().type == DeviceType::cpu) {
        std::copy(data<T>() + i, data<T>() + i + 1, &temp);
      }
  #ifdef CUDA_FOUND
      else {
        gpu::copy(backend_, data<T>() + i, data<T>() + i + 1, &temp);
      }
  #endif
      return temp;
    }
  }

  float get(size_t i) {
    return get<float>(i);
  }

  template <typename T>
  void get(std::vector<T>& v) {
    ABORT_IF(!matchType<T>(type_),
             "Requested type ({}) and underlying type ({}) do not match",
             request<T>(),
             type_);

    v.resize(size());
    if(backend_->getDeviceId().type == DeviceType::cpu) {
      std::copy(data<T>(), data<T>() + size(), v.data());
    }
#ifdef CUDA_FOUND
    else {
      gpu::copy(backend_, data<T>(), data<T>() + size(), v.data());
    }
#endif
  }

  void get(io::Item& item, const std::string& name);

  template <typename T>
  void set(size_t i, T value) {
    if(!matchType<T>(type_)) {
      DISPATCH_BY_TYPE2(type_, set, i, value);
    } else {
      if(backend_->getDeviceId().type == DeviceType::cpu) {
        std::copy(&value, &value + 1, data<T>() + i);
      }
#ifdef CUDA_FOUND
      else {
        gpu::copy(backend_, &value, &value + 1, data<T>() + i);
      }
#endif
    }
  }

  template <typename T>
  void set(const T* begin, const T* end) {
    ABORT_IF(end - begin != shape_.elements(),
             "Vector size ({}) and underlying shape ({}, {}) do not match",
             end - begin,
             std::string(shape_),
             memory_->size());
    matchOrAbort<T>(type_);

    if(backend_->getDeviceId().type == DeviceType::cpu) {
      std::copy(begin, end, data<T>());
    }
#ifdef CUDA_FOUND
    else {
      gpu::copy(backend_, begin, end, data<T>());
    }
#endif
  }

  template <typename T>
  void set(const std::vector<T>& v) {
    set(v.data(), v.data() + v.size());
  }

  // a binary copy with type checking
  void set(const char* begin, const char* end, Type type) {
    ABORT_IF(type_ != type,
             "Tensor type ({}) and data type ({}) do not match",
             type_,
             type);

    size_t dataSize = (end - begin) / sizeOf(type);
    ABORT_IF(size() != dataSize,
             "Tensor size ({}) and mapped size ({}) do not match",
             size(),
             dataSize);

    if(backend_->getDeviceId().type == DeviceType::cpu) {
      std::copy(begin, end, data<char>());
    }
#ifdef CUDA_FOUND
    else {
      gpu::copy(backend_, begin, end, data<char>());
    }
#endif
  }

  void set(const std::vector<char>& v, Type type) {
    set(v.data(), v.data() + v.size(), type);
  }

  void set(const io::Item& item);

  // For single values enable conversion to other numeric formats if possible
  template <typename T>
  void set(T value) {
    if(!matchType<T>(type_)) {
      DISPATCH_BY_TYPE1(type_, setAs, value);
    } else {
      if(backend_->getDeviceId().type == DeviceType::cpu) {
        std::fill(data<T>(), data<T>() + size(), value);
      }
  #ifdef CUDA_FOUND
      else {
        gpu::fill(backend_, data<T>(), data<T>() + size(), value);
      }
  #endif
    }
  }
private: // subroutine for above: helper that accepts any type and casts it to <T>
    template <typename Tas, typename Tval>
    void setAs(Tval value) { set((Tas)value); }
public:

  void setSparse(const std::vector<size_t>& k, const std::vector<float>& v) {
    ABORT_IF(!matchType<float>(type_),
             "Requested type ({}) and underlying type ({}) do not match",
             request<float>(),
             type_);

    if(backend_->getDeviceId().type == DeviceType::cpu) {
      for(size_t i = 0; i < k.size(); ++i)
        data()[k[i]] = v[i];
    }
#ifdef CUDA_FOUND
    else {
      gpu::setSparse(backend_, k, v, data());
    }
#endif
  }

  template <typename T>
  void copyFrom(Tensor in) {
    ABORT_IF(in->shape() != shape_, "Can only copy tensors with equal shapes ({} != {})", in->shape(), shape_);
    ABORT_IF(in->type()  != type_,  "Can only copy tensors with equal types ({} != {})",  in->type(),  type_);
    ABORT_IF(!matchType<T>(type_),
             "Requested type ({}) and underlying type ({}) do not match",
             request<T>(),
             type_);

    if(in->getBackend()->getDeviceId().type == DeviceType::cpu
       && backend_->getDeviceId().type == DeviceType::cpu) {
      std::copy(in->data<T>(), in->data<T>() + in->size(), data<T>());
    }
#ifdef CUDA_FOUND
    else {
      gpu::copy(backend_, in->data<T>(), in->data<T>() + in->size(), data<T>());
    }
#endif
  }

  void copyFrom(Tensor in) {
    DISPATCH_BY_TYPE1(type_, copyFrom, in);
  }

  // Swaps the contents of the current tensor with the argument tensor
  template <typename T>
  void swap(Tensor swapee) {
    ABORT_IF(swapee->shape() != shape_, "Can only swap tensors with equal shapes ({} != {})", swapee->shape(), shape_);
    ABORT_IF(swapee->type()  != type_,  "Can only swap tensors with equal types ({} != {})",  swapee->type(),  type_);
    ABORT_IF(!matchType<T>(type_),
             "Requested type ({}) and underlying type ({}) do not match",
             request<T>(),
             type_);

    // we live on CPUs; just use stdlib
    if(swapee->getBackend()->getDeviceId().type == DeviceType::cpu
       && backend_->getDeviceId().type == DeviceType::cpu) {
      std::swap_ranges(swapee->data<T>(), swapee->data<T>() + swapee->size(), data<T>());
    }
#ifdef CUDA_FOUND
    else {
      if(backend_->getDeviceId() == swapee->getBackend()->getDeviceId()) {
        // we live on the same GPU; do an element-wise swap
        gpu::swap_ranges(backend_, swapee->data<T>(), swapee->data<T>() + swapee->size(), data<T>());
      } else {
        // we live on two different GPUs or devices; go through CPU RAM
        std::vector<T> temp;
        get(temp);
        copyFrom(swapee);
        swapee->set(temp);
      }
    }
#endif
  }

  void swap(Tensor swapee) {
    DISPATCH_BY_TYPE1(type_, swap, swapee);
  }

  template <typename T>
  std::string debug(int precision = 8, int dispCols = 5);

  std::string debug(int precision = 8, int dispCols = 5) {
    DISPATCH_BY_TYPE2(type_, debug, precision, dispCols);
  }

  size_t hash();

};

typedef TensorBase::PtrType Tensor;

template <class TensorType0, class ...TensorTypeRest>
static inline void checkCommonType(TensorType0 first, TensorTypeRest ...rest) {
  std::vector<Tensor> vTensors({first, rest...});
  Type firstType = first->type();
  for(int i = 1; i < vTensors.size(); ++i) {
    ABORT_IF(vTensors[i]->type() != firstType,
             "Type of tensor {} is different from type of tensor 0 ({} != {})",
             i, vTensors[i]->type(), firstType);
  }
}

}  // namespace marian