Remove template parameter from Tensor (pytorch#9939)

Summary:
Pull Request resolved: pytorch#9939

Pull Request resolved: https://github.com/facebookresearch/weakly-supervised-action-detection/pull/13

Pull Request resolved: pytorch/translate#166

Pull Request resolved: pytorch#9125

Closes pytorch#9125

Use inheritance for polymorphism, and remove template parameter
This is to change the templating in call sites, the core implementations will change later

Before Caffe2 Tensor class was compile-time fixed to bind to a particular device/context. With this change, we're making it a runtime property (stored inside the tensor), but preserve the same semantics. For example, one has to specify device type in order to create a Tensor - there are no uninitialized tensors. More specifically the changes are:

1. We added an extra argument *DeviceType* to most of the constructors of the tensor, e.g. (Tensor(DeviceType type)),
2. Semantics of constructor Tensor(const Tensor<SrcContext>& src, ContextForCopy* context); is changed, in this constructor, the second context is passed in to enable us to call the templated Copy function, it could be in a different context as source and target previously, now we'll enforce that the context should have same device type as src, if it is provided.
3. To preserve 'get-or-construct' semantics of Blob, we added specialized getter Blob::GetMutableTensor that verifies both that Blob contains a Tensor and that it's of a correct type
4. Specifically, Tensor type is not default-constructible any more (as we don't have unknown device tensors) and thus some of the code handling STL containers needs to change

Note: Some changes are postponed just to keep this diff a bit smaller. Please see `TODO`s.

Reviewed By: ezyang, houseroad

Differential Revision: D9024330

fbshipit-source-id: e0b8295d2dc6ebe2963383ded5af799ad17164ba

Author: jerryzh168

binaries/benchmark_helper.cc

@@ -160,7 +160,7 @@ void loadInput(
           CAFFE_THROW("Not support GPU on mobile.");
 #endif
         } else {
-          caffe2::TensorCPU* tensor = blob->GetMutable<caffe2::TensorCPU>();
+          caffe2::TensorCPU* tensor = blob->GetMutableTensor(caffe2::CPU);
           CHECK_NOTNULL(tensor);
           tensor->Resize(input_dims);
           if (input_type_list[i] == "uint8_t") {
@@ -197,7 +197,7 @@ void fillInputBlob(
     int protos_size = tensor_kv.second.protos_size();
     caffe2::TensorProto* tensor_proto =
         tensor_kv.second.mutable_protos(iteration % protos_size);
-    caffe2::TensorCPU* tensor = blob->GetMutable<caffe2::TensorCPU>();
+    caffe2::TensorCPU* tensor = blob->GetMutableTensor(caffe2::CPU);
     tensor->Resize(std::vector<caffe2::TIndex>());
     if (tensor_proto->data_type() == caffe2::TensorProto::STRING) {
       (tensor->mutable_data<std::string>())[0] = tensor_proto->string_data(0);
@@ -290,7 +290,7 @@ void writeOutput(
 #endif
         } else {
           writeTextOutput<caffe2::CPUContext, caffe2::TensorCPU>(
-              workspace->GetBlob(name)->GetMutable<caffe2::TensorCPU>(),
+              workspace->GetBlob(name)->GetMutableTensor(caffe2::CPU),
               output_prefix,
               name);
         }

binaries/benchmark_helper.h

@@ -35,7 +35,7 @@ void writeTextOutput(
     const string& output_prefix,
     const string& name) {
   string output_name = output_prefix + "/" + name + ".txt";
-  caffe2::TensorSerializer<ContextType> ser;
+  caffe2::TensorSerializer ser;
   caffe2::BlobProto blob_proto;
   ser.Serialize(
       *tensor, output_name, blob_proto.mutable_tensor(), 0, tensor->size());

binaries/core_overhead_benchmark.cc

@@ -139,7 +139,7 @@ BENCHMARK(BM_cudaStreamWaitEventThenStreamSynchronize);
 
 static void BM_CudaPointerAffinity(benchmark::State& state) {
   CAFFE2_SKIP_IF_NO_GPU;
-  TensorCUDA tensor(vector<TIndex>{1, 2, 3, 4});
+  Tensor tensor(vector<TIndex>{1, 2, 3, 4}, CUDA);
   float* ptr = tensor.mutable_data<float>();
   while (state.KeepRunning()) {
     volatile int id = GetGPUIDForPointer(ptr);
@@ -198,7 +198,7 @@ static void BM_RawAllocDeallocCPU(benchmark::State& state) {
 BENCHMARK(BM_RawAllocDeallocCPU);
 
 static void BM_TensorAllocDeallocCPU(benchmark::State& state) {
-  Tensor<CPUContext> tensor;
+  Tensor tensor(CPU);
   // small allocation
   tensor.Resize(32, 32);
   while (state.KeepRunning()) {
@@ -210,7 +210,7 @@ BENCHMARK(BM_TensorAllocDeallocCPU);
 
 static void BM_TensorAllocDeallocCUDA(benchmark::State& state) {
   CAFFE2_SKIP_IF_NO_GPU;
-  Tensor<CUDAContext> tensor;
+  Tensor tensor(CUDA);
   // small allocation
   tensor.Resize(32, 32);
   while (state.KeepRunning()) {

binaries/print_core_object_sizes.cc

@@ -28,8 +28,7 @@
 
 int main(int /* unused */, char** /* unused */) {
   PRINT_SIZE(caffe2::Blob);
-  PRINT_SIZE(caffe2::Tensor<caffe2::CPUContext>);
-  PRINT_SIZE(caffe2::Tensor<caffe2::CUDAContext>);
+  PRINT_SIZE(caffe2::Tensor);
   PRINT_SIZE(caffe2::CPUContext);
   PRINT_SIZE(caffe2::CUDAContext);
   PRINT_SIZE(caffe2::OperatorBase);

binaries/speed_benchmark.cc

@@ -136,7 +136,7 @@ int main(int argc, char** argv) {
         if (blob == nullptr) {
           blob = workspace->CreateBlob(input_names[i]);
         }
-        caffe2::TensorCPU* tensor = blob->GetMutable<caffe2::TensorCPU>();
+        caffe2::TensorCPU* tensor = blob->GetMutableTensor(caffe2::CPU);
         CHECK_NOTNULL(tensor);
         tensor->Resize(input_dims);
         if (input_type_list[i] == "uint8_t") {

caffe2/contrib/aten/aten_op_template.h

@@ -54,11 +54,11 @@ class ATenOp : public Operator<Context> {
     #undef DEFINE_CASE
   }
 
-  at::Type & typeFor(const Tensor<Context> & ten) {
+  at::Type& typeFor(const Tensor& ten) {
     return at::getType(backend(), atScalarTypeFor(ten.meta()));
   }
-  at::Tensor tensorWrapping(const Tensor<Context>& ten_) {
-    auto& ten = const_cast<Tensor<Context>&>(ten_);
+  at::Tensor tensorWrapping(const Tensor& ten_) {
+    auto& ten = const_cast<Tensor&>(ten_);
     return typeFor(ten).tensorFromBlob(ten.raw_mutable_data(), ten.dims());
   }
 
@@ -88,7 +88,7 @@ class ATenOp : public Operator<Context> {
     }
     CAFFE_THROW("Unknown type meta"); // TODO: improve error message...
   }
-  void assignTo(Tensor<Context> * dst, const at::Tensor & src_) {
+  void assignTo(Tensor* dst, const at::Tensor& src_) {
     at::Tensor src = src_.contiguous();
     auto at_sizes = src.sizes();
     std::vector<int64_t> dims(at_sizes.begin(),at_sizes.end());
@@ -121,7 +121,7 @@ class ATenOp : public Operator<Context> {
     return s.toLong();
   }
 
-  void assignTo(Tensor<Context> * dst, at::Type & inferred_type, at::Scalar scalar) {
+  void assignTo(Tensor* dst, at::Type& inferred_type, at::Scalar scalar) {
     switch(inferred_type.scalarType()) {
       #define DEFINE_CASE(ctype,aten_name,native) \
         case at::k##aten_name: { \
@@ -134,8 +134,8 @@ class ATenOp : public Operator<Context> {
         CAFFE_THROW("Unknown ATen Type");
     }
   }
-  template<typename T>
-  void assignToValue(Tensor<Context> * dst, T v) {
+  template <typename T>
+  void assignToValue(Tensor* dst, T v) {
     dst->Resize(std::vector<TIndex>());
     math::Set(1, v, dst->template mutable_data<T>(), &context_);
   }

caffe2/contrib/gloo/common.cc

@@ -12,7 +12,7 @@ namespace caffe2 {
 namespace gloo {
 
 void signalFailure(Blob* status_blob, std::exception& /* unused */) {
-  auto* res = status_blob->GetMutable<TensorCPU>();
+  auto* res = status_blob->GetMutableTensor(CPU);
   res->Resize(1);
   res->template mutable_data<int32_t>()[0] = 1;
 }

caffe2/contrib/nccl/cuda_nccl_op_gpu.cc

@@ -17,17 +17,17 @@ nccl::NCCLExecution getNCCLElements(
   ex.elements.resize(op->InputSize());
   for (auto i = 0; i < op->InputSize(); ++i) {
     auto& el = ex.elements[i];
-    el.src = &(op->Input<TensorCUDA>(i));
+    el.src = &(op->Input<Tensor>(i, CUDA));
     if (op->OutputSize() == 1) {
       // Reduce op
       if (i == ex.root) {
-        el.dst = op->Output<TensorCUDA>(0);
+        el.dst = op->Output<Tensor>(0, CUDA);
       }
     } else if (i < op->OutputSize()) {
-      el.dst = op->Output<TensorCUDA>(i);
+      el.dst = op->Output<Tensor>(i, CUDA);
     }
     // TODO - expensive (>1ms) - cache these.
-    el.device = GetGPUIDForPointer(op->Input<TensorCUDA>(i).raw_data());
+    el.device = GetGPUIDForPointer(op->Input<Tensor>(i, CUDA).raw_data());
   }
 
   return ex;
@@ -38,7 +38,7 @@ namespace {
 template <typename T>
 bool AllInputsAre(OperatorBase* op) {
   for (auto i = 0; i < op->InputSize(); ++i) {
-    if (op->Input<TensorCUDA>(i).IsType<T>()) {
+    if (op->Input<Tensor>(i, CUDA).IsType<T>()) {
       continue;
     } else {
       return false;

caffe2/contrib/nervana/nervana_fc_op_gpu_test.cc

@@ -22,7 +22,7 @@ static void AddConstInput(const std::vector<int>& shape, const float value,
   option.set_device_type(CUDA);
   CUDAContext context(option);
   Blob* blob = ws->CreateBlob(name);
-  auto* tensor = blob->GetMutable<Tensor<CUDAContext>>();
+  auto* tensor = blob->GetMutableTensor(CUDA);
   tensor->Resize(shape);
   math::Set<float, CUDAContext>(tensor->size(), value,
                                 tensor->mutable_data<float>(),
@@ -54,8 +54,8 @@ TEST(NervanaFullyConnectedTest, Test) {
   EXPECT_TRUE(op->Run());
   Blob* Yblob = ws.GetBlob("Y");
   EXPECT_NE(nullptr, Yblob);
-  auto& Y = Yblob->Get<Tensor<CUDAContext>>();
-  TensorCPU Y_cpu(Y);
+  auto& Y = Yblob->Get<Tensor>();
+  Tensor Y_cpu(Y, CPU);
   EXPECT_EQ(Y.size(), 5 * 6);
   for (int i = 0; i < Y.size(); ++i) {
     CHECK_LT(Y_cpu.data<float>()[i], 10.11);

caffe2/contrib/warpctc/ctc_op.h

@@ -47,26 +47,26 @@ class CTCOp final : public Operator<Context> {
     const auto& inputs = Input(INPUTS);
     const auto minibatchSize = inputs.dim(1);
     const auto alphabetSize = inputs.dim(2);
-    const auto& labels = OperatorBase::template Input<TensorCPU>(LABELS);
+    const auto& labels = OperatorBase::template Input<Tensor>(LABELS, CPU);
     const auto& labelLengths =
-        OperatorBase::template Input<TensorCPU>(LABEL_LENGTHS);
+        OperatorBase::template Input<Tensor>(LABEL_LENGTHS, CPU);
     const auto& inputLengths =
-        OperatorBase::template Input<TensorCPU>(INPUT_LENGTHS);
+        OperatorBase::template Input<Tensor>(INPUT_LENGTHS, CPU);
 
     // outputs
-    Tensor<Context>* gradients = nullptr;
+    Tensor* gradients = nullptr;
     TensorCPU* costs;
-    Tensor<Context>* workspace;
+    Tensor* workspace;
     if (!is_test_) {
       // [grads, costs, workspace] to maintain backward compatibility
       gradients = Output(0);
       gradients->ResizeLike(inputs);
-      costs = OperatorBase::template Output<TensorCPU>(1);
+      costs = OperatorBase::template Output<Tensor>(1, CPU);
       costs->ResizeLike(labelLengths);
       workspace = Output(2);
     } else {
       // [costs, workspace]
-      costs = OperatorBase::template Output<TensorCPU>(0);
+      costs = OperatorBase::template Output<Tensor>(0, CPU);
       costs->ResizeLike(labelLengths);
       workspace = Output(1);
     }

caffe2/core/allocator.cc

@@ -26,7 +26,7 @@ void SetCPUAllocator(CPUAllocator* alloc) {
   g_cpu_allocator.reset(alloc);
 }
 
-MemoryAllocationReporter CPUContext::reporter_;
+MemoryAllocationReporter CPUStaticContext::reporter_;
 
 void MemoryAllocationReporter::New(void* ptr, size_t nbytes) {
   std::lock_guard<std::mutex> guard(mutex_);

caffe2/core/blob.h

@@ -9,8 +9,9 @@
 
 #include "caffe2/core/blob_serializer_base.h"
 #include "caffe2/core/common.h"
-#include "caffe2/core/typeid.h"
 #include "caffe2/core/logging.h"
+#include "caffe2/core/tensor.h"
+#include "caffe2/core/typeid.h"
 #include "caffe2/proto/caffe2.pb.h"
 
 namespace caffe2 {
@@ -60,6 +61,20 @@ class Blob {
   template <class T>
   bool IsType() const { return meta_.Match<T>(); }
 
+  // TODO(jerryzh): Remove template
+  template <class T>
+  bool IsType(DeviceType device_type) const {
+    static_assert(
+        std::is_same<T, Tensor>::value,
+        "IsType(DeviceType) only available on "
+        "Tensor types.");
+    auto* tensor = static_cast<Tensor*>(pointer_);
+    if (tensor && tensor->GetDeviceType() == device_type) {
+      return true;
+    }
+    return false;
+  }
+
   /**
    * Returns the meta info of the blob.
    */
@@ -74,6 +89,7 @@ class Blob {
    * @brief Gets the const reference of the stored object. The code checks if
    * the stored object is of the desired type.
    */
+  // TODO(jerryzh): add a Get(DeviceType) function?
   template <class T>
   const T& Get() const {
     CAFFE_ENFORCE(
@@ -123,6 +139,17 @@ class Blob {
     }
   }
 
+  inline Tensor* GetMutableTensor(DeviceType device_type) {
+    if (IsType<Tensor>() &&
+        static_cast<Tensor*>(pointer_)->GetDeviceType() == device_type) {
+      return static_cast<Tensor*>(pointer_);
+    } else {
+      VLOG(1) << "Create new mutable object " << TypeMeta::TypeName<Tensor>()
+              << " DeviceType:" << device_type;
+      return Reset<Tensor>(new Tensor(device_type));
+    }
+  }
+
   /**
    * Sets the underlying object to the allocated one. The Blob then takes over
    * the ownership of the passed in pointer. If there is already an object in