Remove almost all usage of Blob in EagerBlobObject

oneflow/api/python/functional/tensor_api.cpp

@@ -280,8 +280,6 @@ class LocalTensorSharedNumpyDataFunctor {
     JUST(tensor_impl->InitEagerBlobObject(NewLocalDepObject()));
     const auto& stream = GetDefaultStreamByDevice(device);
     JUST(tensor_impl->eager_blob_object())->set_last_used_stream(stream);
-    JUST(JUST(tensor_impl->eager_blob_object())->TryInitBlob());
-    JUST(tensor_impl->eager_blob_object())->mut_blob()->reset_dptr(static_cast<char*>(data_ptr));
     std::shared_ptr<Tensor> out(new MirroredTensor(tensor_impl));
     return out;
   }

oneflow/api/python/utils/tensor_utils.cpp

@@ -100,9 +100,9 @@ MaybeGetTensorBufferShapesAndDTypes(const std::shared_ptr<Tensor>& t) {
   }));
   JUST(btb->WaitUntilCntEqualZero(VirtualMachine::GetPredicatorNoMoreInstructionsFinished()));
 
-  const Blob& blob = JUST(tensor->eager_blob_object())->blob();
-  const Shape& blob_shape = blob.static_shape();
-  const auto* tensor_buffer_ptr = blob.dptr<TensorBuffer>();
+  const auto& eager_blob_object = JUST(tensor->eager_blob_object());
+  const Shape& blob_shape = eager_blob_object->shape();
+  const auto* tensor_buffer_ptr = eager_blob_object->dptr<TensorBuffer>();
   for (int64_t i = 0; i < blob_shape.elem_cnt(); ++i) {
     const TensorBuffer* tensor_buffer = tensor_buffer_ptr + i;
     shapes.emplace_back(tensor_buffer->shape());

oneflow/core/common/shape.cpp

@@ -60,34 +60,31 @@ int64_t ShiftNegativeAxis(int64_t axis, const int64_t num_axes) {
   return axis;
 }
 
-Shape::Shape(const std::initializer_list<int64_t>& dim_vec) : dim_vec_(dim_vec) { UpdateElemCnt(); }
-Shape::Shape(const DimVector& dim_vec) : dim_vec_(dim_vec) { UpdateElemCnt(); }
-Shape::Shape(DimVector&& dim_vec) : dim_vec_(std::move(dim_vec)) { UpdateElemCnt(); }
-Shape::Shape(const ShapeProto& shape_proto) {
+Shape::Shape(const std::initializer_list<int64_t>& dim_vec)
+    : dim_vec_(dim_vec), is_initialized_(true) {}
+Shape::Shape(const DimVector& dim_vec) : dim_vec_(dim_vec), is_initialized_(true) {}
+Shape::Shape(DimVector&& dim_vec) : dim_vec_(std::move(dim_vec)), is_initialized_(true) {}
+Shape::Shape(const ShapeProto& shape_proto) : is_initialized_(true) {
   dim_vec_.assign(shape_proto.dim().begin(), shape_proto.dim().end());
-  UpdateElemCnt();
 }
-Shape::Shape(const cfg::ShapeProto& shape_proto) {
+Shape::Shape(const cfg::ShapeProto& shape_proto) : is_initialized_(true) {
   dim_vec_.assign(shape_proto.dim().begin(), shape_proto.dim().end());
-  UpdateElemCnt();
 }
 
 Shape& Shape::operator=(const Shape& shape) {
   dim_vec_ = shape.dim_vec_;
-  UpdateElemCnt();
+  is_initialized_ = shape.is_initialized_;
   return *this;
 }
 
 Shape& Shape::assign(const DimVector& dim_vec) {
   dim_vec_ = dim_vec;
-  UpdateElemCnt();
   return *this;
 }
 
 Shape& Shape::CheckNumAxesIdenticalAndAssign(const ShapeView& shape_view) {
   CHECK_EQ(NumAxes(), shape_view.NumAxes());
   std::copy(shape_view.ptr(), shape_view.ptr() + shape_view.NumAxes(), dim_vec_.data());
-  UpdateElemCnt();
   return *this;
 }
 
@@ -97,7 +94,6 @@ Shape& Shape::LeftOnesExtendedAssign(const ShapeView& shape_view) {
   FOR_RANGE(int, i, 0, left_ones_size) { dim_vec_.at(i) = 1LL; }
   std::copy(shape_view.ptr(), shape_view.ptr() + shape_view.NumAxes(),
             dim_vec_.data() + left_ones_size);
-  UpdateElemCnt();
   return *this;
 }
 
@@ -133,7 +129,6 @@ void Shape::Set(int64_t index, int64_t val) {
   CHECK_LT(index, this->NumAxes()) << " Shape: " << DebugStr() << " visit index: " << index
                                    << " > num_axes: " << this->NumAxes();
   dim_vec_.at(index) = val;
-  UpdateElemCnt();
 }
 
 int64_t Shape::Count(int64_t begin_axis, int64_t end_axis) const {
@@ -146,12 +141,6 @@ int64_t Shape::Count(int64_t begin_axis, int64_t end_axis) const {
 
 int64_t Shape::Count(int64_t begin_axis) const { return Count(begin_axis, NumAxes()); }
 
-void Shape::UpdateElemCnt() {
-  int64_t elem_cnt = 1;
-  for (int64_t s : dim_vec_) { elem_cnt *= s; }
-  elem_cnt_ = elem_cnt;
-}
-
 std::ostream& operator<<(std::ostream& out, const Shape& shape) {
   out << shape.DebugStr();
   return out;

oneflow/core/common/shape.h

@@ -17,6 +17,7 @@ limitations under the License.
 #define ONEFLOW_CORE_COMMON_SHAPE_H_
 
 #include "oneflow/core/common/shape.pb.h"
+#include "oneflow/core/common/shape_view.h"
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/common/maybe.h"
 #include "oneflow/core/common/shape_vec.h"
@@ -33,7 +34,7 @@ class ShapeProto;
 class Shape final {
  public:
   // OF_DISALLOW_COPY_AND_MOVE(Shape);
-  Shape() = default;
+  Shape() : is_initialized_(false) {}
   explicit Shape(const DimVector& dim_vec);
   explicit Shape(DimVector&& dim_vec);
   explicit Shape(const ShapeProto& shape_proto);
@@ -56,9 +57,12 @@ class Shape final {
   void SerializeWithTextFormat(StreamT& out_stream) const;
 
   // Getters and Setters
-  bool is_initialized() const { return elem_cnt_.has_value(); }
+  bool is_initialized() const { return is_initialized_; }
   const DimVector& dim_vec() const { return dim_vec_; }
-  int64_t elem_cnt() const { return CHECK_JUST(elem_cnt_); }
+  DimVector& dim_vec() { return dim_vec_; }
+  int64_t elem_cnt() const {
+    return std::accumulate(dim_vec_.begin(), dim_vec_.end(), 1, std::multiplies<int>());
+  }
   int64_t At(int64_t index) const;
   void Set(int64_t index, int64_t val);
   int64_t NumAxes() const {
@@ -78,11 +82,13 @@ class Shape final {
 
   Maybe<Shape> Slice(int64_t start_dim, int64_t end_dim) const;
 
- private:
-  void UpdateElemCnt();
+  ShapeView ToShapeView() const { return ShapeView(dim_vec_.data(), dim_vec_.size()); }
 
+  MutShapeView ToMutShapeView() { return MutShapeView(dim_vec_.data(), dim_vec_.size()); }
+
+ private:
   DimVector dim_vec_;
-  Optional<int64_t> elem_cnt_;
+  bool is_initialized_;
 };
 
 int64_t ShiftNegativeAxis(int64_t axis, const int64_t num_axes);

oneflow/core/common/util.h

@@ -181,6 +181,9 @@ inline double GetCurTime() {
 const size_t kHostAlignSize = 64;
 const size_t kCudaAlignSize = 512;
 const size_t kCudaMemAllocAlignSize = 512;
+const int32_t kBlobBodyAlignSize = 512;
+const int32_t kBlobHeaderAlignSize = 64;
+
 inline size_t RoundUp(size_t n, size_t val) { return (n + val - 1) / val * val; }
 
 inline size_t GetCudaAlignedSize(size_t size) { return RoundUp(size, kCudaAlignSize); }

oneflow/core/eager/blob_instruction_type.cpp

@@ -19,42 +19,18 @@ limitations under the License.
 #include "oneflow/core/vm/instruction.h"
 #include "oneflow/core/vm/instruction_type.h"
 #include "oneflow/core/eager/blob_instruction_type.h"
-#include "oneflow/core/eager/blob_object.h"
 #include "oneflow/core/vm/control_stream_type.h"
 #include "oneflow/core/vm/stream.h"
 #include "oneflow/core/device/cuda_util.h"
 #include "oneflow/core/register/register_manager.h"
-#include "oneflow/core/eager/lazy_ref_blob_object.h"
 #include "oneflow/core/operator/operator.h"
 #include "oneflow/core/vm/access_blob_arg_cb_phy_instr_operand.h"
 #include "oneflow/core/register/ofblob.h"
 #include "oneflow/core/eager/eager_blob_object.h"
-#include "oneflow/core/vm/tensor_view_operand.h"
 
 namespace oneflow {
 namespace vm {
 
-void TensorViewInstructionType::Compute(vm::Instruction* instruction) const {
-  return ComputeInstrMsg(instruction->instr_msg());
-}
-
-void TensorViewInstructionType::ComputeInFuseMode(vm::InstructionMsg* instr_msg) const {
-  return ComputeInstrMsg(*instr_msg);
-}
-
-void TensorViewInstructionType::ComputeInstrMsg(const vm::InstructionMsg& instr_msg) const {
-  const auto& phy_instr_operand = instr_msg.phy_instr_operand();
-  CHECK(static_cast<bool>(phy_instr_operand));
-  const auto* ptr = dynamic_cast<const vm::TensorViewOperand*>(phy_instr_operand.get());
-  CHECK_NOTNULL(ptr);
-  DeviceCtx* device_ctx = instr_msg.phy_instr_stream()->device_ctx().get();
-  OfBlob input_ofblob(device_ctx->stream(), ptr->eager_blob_object()->mut_blob());
-  OfBlob view_ofblob(device_ctx->stream(), ptr->view_eager_blob_object()->mut_blob());
-
-  void* input_ptr = input_ofblob.mut_blob()->mut_raw_dptr();
-  view_ofblob.mut_blob()->reset_dptr(static_cast<char*>(input_ptr));
-}
-
 void AccessBlobByCallbackInstructionType::Compute(vm::Instruction* instruction) const {
   return ComputeInstrMsg(instruction->instr_msg());
 }
@@ -71,7 +47,8 @@ void AccessBlobByCallbackInstructionType::ComputeInstrMsg(
       dynamic_cast<const vm::AccessBlobArgCbPhyInstrOperand*>(phy_instr_operand.get());
   CHECK_NOTNULL(ptr);
   DeviceCtx* device_ctx = instr_msg.phy_instr_stream()->device_ctx().get();
-  OfBlob ofblob(device_ctx->stream(), ptr->eager_blob_object()->mut_blob());
+  auto* blob = ptr->eager_blob_object()->blob();
+  OfBlob ofblob(device_ctx->stream(), blob);
   ptr->callback()(reinterpret_cast<uint64_t>(&ofblob));
 }
 

oneflow/core/eager/blob_instruction_type.h

@@ -19,18 +19,6 @@ limitations under the License.
 namespace oneflow {
 namespace vm {
 
-class TensorViewInstructionType : public vm::InstructionType {
- public:
-  TensorViewInstructionType() = default;
-  ~TensorViewInstructionType() override = default;
-
-  void Compute(vm::Instruction* instruction) const override;
-  void ComputeInFuseMode(vm::InstructionMsg* instr_msg) const override;
-
- private:
-  void ComputeInstrMsg(const vm::InstructionMsg& instr_msg) const;
-};
-
 class AccessBlobByCallbackInstructionType : public vm::InstructionType {
  public:
   AccessBlobByCallbackInstructionType() = default;

oneflow/core/eager/cpu_blob_instruction_type.cpp

@@ -29,14 +29,5 @@ class CpuAccessBlobByCallbackInstructionType final : public AccessBlobByCallback
 COMMAND(vm::RegisterInstructionType<CpuAccessBlobByCallbackInstructionType>(
     "cpu.AccessBlobByCallback"));
 
-class CpuTensorViewInstructionType final : public TensorViewInstructionType {
- public:
-  CpuTensorViewInstructionType() = default;
-  ~CpuTensorViewInstructionType() override = default;
-
-  using stream_type = vm::CpuStreamType;
-};
-COMMAND(vm::RegisterInstructionType<CpuTensorViewInstructionType>("cpu.TensorView"));
-
 }  // namespace vm
 }  // namespace oneflow

oneflow/core/eager/cpu_opkernel_instruction_type.cpp

@@ -15,7 +15,6 @@ limitations under the License.
 */
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/job/job_desc.h"
-#include "oneflow/core/eager/blob_object.h"
 #include "oneflow/core/eager/opkernel_instruction_type.h"
 #include "oneflow/core/vm/stream.h"
 #include "oneflow/core/vm/cpu_stream_type.h"