Symbolic local tensor meta

oneflow/api/python/functional/tensor_api.cpp

@@ -266,7 +266,7 @@ class LocalTensorSharedNumpyDataFunctor {
       }
       stride_val /= element_size_in_bytes;
     }
-    auto tensor_meta = std::make_shared<LocalTensorMeta>(shape, strides, data_type, device, 0);
+    auto tensor_meta = SymbolOf(LocalTensorMeta(shape, strides, data_type, device, 0));
 
     // Build TensorBuffer
     const auto& Free = [array](char* dptr) {
@@ -286,12 +286,12 @@ class LocalTensorSharedNumpyDataFunctor {
     auto tensor_storage = std::make_shared<TensorStorage>(tensor_data);
 
     // Build Tensor
-    auto tensor_impl = std::make_shared<EagerLocalTensorImpl>(tensor_meta, tensor_storage,
+    auto tensor_impl = std::make_shared<EagerLocalTensorImpl>(tensor_storage,
                                                               /*requires_grad=*/false,
                                                               /*ls_leaf=*/true);
 
     // Init blob
-    JUST(tensor_impl->InitEagerBlobObject(NewLocalDepObject()));
+    JUST(tensor_impl->InitEagerBlobObject(tensor_meta, NewLocalDepObject()));
     const auto& stream = JUST(GetDefaultStreamByDevice(device));
     const auto& eager_blob_object = JUST(tensor_impl->eager_blob_object());
     JUST(eager_blob_object->init_producer_stream(stream));

oneflow/core/framework/tensor_desc.cpp → oneflow/core/common/tensor_desc.cpp

@@ -13,7 +13,8 @@ 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 "oneflow/core/framework/tensor_desc.h"
+#include "oneflow/core/common/tensor_desc.h"
+#include "oneflow/core/register/blob_desc.pb.h"
 
 namespace oneflow {
 

oneflow/core/framework/tensor_desc.h → oneflow/core/common/tensor_desc.h

@@ -13,16 +13,18 @@ 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.
 */
-#ifndef ONEFLOW_CORE_FRAMEWORK_TENSOR_DESC_H_
-#define ONEFLOW_CORE_FRAMEWORK_TENSOR_DESC_H_
+#ifndef ONEFLOW_CORE_COMMON_TENSOR_DESC_H_
+#define ONEFLOW_CORE_COMMON_TENSOR_DESC_H_
 
 #include "oneflow/core/common/util.h"
-#include "oneflow/core/register/blob_desc.pb.h"
 #include "oneflow/core/common/shape.h"
 #include "oneflow/core/common/stride.h"
+#include "oneflow/core/common/data_type.pb.h"
 
 namespace oneflow {
 
+class BlobDescProto;
+
 namespace user_op {
 
 class TensorDesc {
@@ -77,4 +79,4 @@ class NaiveTensorDesc final : public TensorDesc {
 
 }  // namespace oneflow
 
-#endif  // ONEFLOW_CORE_FRAMEWORK_TENSOR_DESC_H_
+#endif  // ONEFLOW_CORE_COMMON_TENSOR_DESC_H_

oneflow/core/framework/tensor_meta.cpp → oneflow/core/common/tensor_meta.cpp

@@ -13,13 +13,36 @@ 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 "oneflow/core/framework/tensor_meta.h"
+#include "oneflow/core/common/tensor_meta.h"
 #include "oneflow/core/common/stride.h"
 #include "oneflow/core/framework/device.h"
 
 namespace oneflow {
 namespace one {
 
+MutTensorMeta::MutTensorMeta()
+    : TensorMeta(std::make_shared<const Shape>(), std::make_shared<const Stride>(),
+                 kInvalidDataType) {}
+
+MutTensorMeta::MutTensorMeta(const std::shared_ptr<const Shape>& shape, DataType dtype)
+    : TensorMeta(shape, std::make_shared<const Stride>(*shape), dtype) {}
+
+MutTensorMeta::MutTensorMeta(const std::shared_ptr<const Shape>& shape,
+                             const std::shared_ptr<const Stride>& stride, DataType dtype)
+    : TensorMeta(shape, stride, dtype) {}
+
+bool MutTensorMeta::operator==(const MutTensorMeta& other) const {
+  // It's correct to ignore is_dynamic_ field.
+  return *this->shape_ptr() == *other.shape_ptr() && this->dtype() == other.dtype()
+         && this->stride() == other.stride();
+}
+
+size_t MutTensorMeta::CalcHashValue() const {
+  // It's correct to ignore is_dynamic_ field.
+  return std::hash<Shape>()(*shape_ptr()) ^ std::hash<DataType>()(dtype())
+         ^ std::hash<Stride>()(stride());
+}
+
 LocalTensorMeta::LocalTensorMeta()
     : TensorMeta(std::make_shared<const Shape>(), std::make_shared<const Stride>(),
                  DataType::kInvalidDataType),
@@ -50,6 +73,36 @@ size_t LocalTensorMeta::CalcHashValue() const {
          ^ std::hash<Device>()(*device()) ^ std::hash<Stride>()(stride()) ^ storage_offset();
 }
 
+MutLocalTensorMeta::MutLocalTensorMeta()
+    : MutTensorMeta(std::make_shared<const Shape>(), std::make_shared<const Stride>(),
+                    kInvalidDataType),
+      device_(Symbol<Device>()),
+      storage_offset_(0) {}
+
+MutLocalTensorMeta::MutLocalTensorMeta(const std::shared_ptr<const Shape>& shape, DataType dtype,
+                                       Symbol<Device> device)
+    : MutTensorMeta(shape, std::make_shared<const Stride>(*shape), dtype),
+      device_(device),
+      storage_offset_(0) {}
+
+MutLocalTensorMeta::MutLocalTensorMeta(const std::shared_ptr<const Shape>& shape,
+                                       const std::shared_ptr<const Stride>& stride, DataType dtype,
+                                       Symbol<Device> device, int64_t storage_offset)
+    : MutTensorMeta(shape, stride, dtype), device_(device), storage_offset_(storage_offset) {}
+
+bool MutLocalTensorMeta::operator==(const MutLocalTensorMeta& other) const {
+  // It's correct to ignore is_dynamic_ field.
+  return *this->shape_ptr() == *other.shape_ptr() && this->dtype() == other.dtype()
+         && *this->device() == *other.device() && this->stride() == other.stride()
+         && this->storage_offset() == other.storage_offset();
+}
+
+size_t MutLocalTensorMeta::CalcHashValue() const {
+  // It's correct to ignore is_dynamic_ field.
+  return std::hash<Shape>()(*shape_ptr()) ^ std::hash<DataType>()(dtype())
+         ^ std::hash<Device>()(*device()) ^ std::hash<Stride>()(stride()) ^ storage_offset();
+}
+
 bool GlobalTensorMeta::operator==(const GlobalTensorMeta& other) const {
   // It's correct to ignore is_dynamic_ field.
   return *this->shape_ptr() == *other.shape_ptr() && this->dtype() == other.dtype()

oneflow/core/framework/tensor_meta.h → oneflow/core/common/tensor_meta.h

@@ -13,11 +13,11 @@ 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.
 */
-#ifndef ONEFLOW_FRAMEWORK_TENSOR_META_H_
-#define ONEFLOW_FRAMEWORK_TENSOR_META_H_
+#ifndef ONEFLOW_COMMON_TENSOR_META_H_
+#define ONEFLOW_COMMON_TENSOR_META_H_
 
 #include <memory>
-#include "oneflow/core/framework/tensor_desc.h"
+#include "oneflow/core/common/tensor_desc.h"
 #include "oneflow/core/common/symbol.h"
 
 namespace oneflow {
@@ -60,15 +60,23 @@ class TensorMeta : public user_op::TensorDesc {
   bool is_dynamic() const override { return is_dynamic_; }
   bool is_contiguous() const { return IsContiguous(shape(), *stride_); }
 
-  void set_shape(const std::shared_ptr<const Shape>& val) { shape_ = val; }
-  Shape* mut_shape() override { return const_cast<Shape*>(shape_.get()); }
-  void set_stride(const std::shared_ptr<const Stride>& val) { stride_ = val; }
-  Stride* mut_stride() override { return const_cast<Stride*>(stride_.get()); }
-  DataType* mut_dtype() { return &data_type_; }
-  void set_dtype(DataType data_type) { data_type_ = data_type; }
-  DataType* mut_data_type() override { return &data_type_; }
-  bool* mut_is_dynamic() override { return &is_dynamic_; }
-  void set_is_dynamic(bool val) override { is_dynamic_ = val; }
+  virtual Shape* mut_shape() override {
+    PRINT_BUG_PROMPT_AND_ABORT();
+    return nullptr;
+  }
+  virtual Stride* mut_stride() override {
+    PRINT_BUG_PROMPT_AND_ABORT();
+    return nullptr;
+  }
+  virtual DataType* mut_data_type() override {
+    PRINT_BUG_PROMPT_AND_ABORT();
+    return nullptr;
+  }
+  virtual bool* mut_is_dynamic() override {
+    PRINT_BUG_PROMPT_AND_ABORT();
+    return nullptr;
+  }
+  virtual void set_is_dynamic(bool val) override { PRINT_BUG_PROMPT_AND_ABORT(); }
 
  protected:
   TensorMeta& operator=(const TensorMeta& other) {
@@ -79,13 +87,39 @@ class TensorMeta : public user_op::TensorDesc {
     return *this;
   }
 
- private:
   std::shared_ptr<const Shape> shape_;
   std::shared_ptr<const Stride> stride_;
   DataType data_type_;
   bool is_dynamic_;
 };
 
+class MutTensorMeta : public TensorMeta {
+ public:
+  // uninitialized MutTensorMeta.
+  MutTensorMeta();
+  MutTensorMeta(const MutTensorMeta&) = default;
+  MutTensorMeta(const std::shared_ptr<const Shape>& shape, DataType dtype);
+  MutTensorMeta(const std::shared_ptr<const Shape>& shape,
+                const std::shared_ptr<const Stride>& stride, DataType dtype);
+  virtual ~MutTensorMeta() = default;
+
+  Shape* mut_shape() override { return const_cast<Shape*>(shape_.get()); }
+  Stride* mut_stride() override { return const_cast<Stride*>(stride_.get()); }
+  DataType* mut_data_type() override { return &data_type_; }
+  bool* mut_is_dynamic() override { return &is_dynamic_; }
+  void set_is_dynamic(bool val) override { is_dynamic_ = val; }
+
+  void set_shape(const std::shared_ptr<const Shape>& val) { shape_ = val; }
+  void set_stride(const std::shared_ptr<const Stride>& val) { stride_ = val; }
+  DataType* mut_dtype() { return &data_type_; }
+  void set_dtype(DataType data_type) { data_type_ = data_type; }
+
+  bool operator==(const MutTensorMeta& other) const;
+  size_t CalcHashValue() const;
+
+  MutTensorMeta& operator=(const MutTensorMeta& other) = default;
+};
+
 class LocalTensorMeta : public TensorMeta {
  public:
   // uninitialized LocalTensorMeta.
@@ -100,13 +134,38 @@ class LocalTensorMeta : public TensorMeta {
   const Symbol<Device>& device() const { return device_; }
   int64_t storage_offset() const { return storage_offset_; }
 
+  bool operator==(const LocalTensorMeta& other) const;
+  size_t CalcHashValue() const;
+
+  LocalTensorMeta& operator=(const LocalTensorMeta& other) = default;
+
+ private:
+  Symbol<Device> device_;
+  int64_t storage_offset_;
+};
+
+class MutLocalTensorMeta : public MutTensorMeta {
+ public:
+  // uninitialized MutLocalTensorMeta.
+  MutLocalTensorMeta();
+  MutLocalTensorMeta(const MutLocalTensorMeta&) = default;
+  MutLocalTensorMeta(const std::shared_ptr<const Shape>& shape, DataType dtype,
+                     Symbol<Device> device);
+  MutLocalTensorMeta(const std::shared_ptr<const Shape>& shape,
+                     const std::shared_ptr<const Stride>& stride, DataType dtype,
+                     Symbol<Device> device, int64_t storage_offset);
+  virtual ~MutLocalTensorMeta() = default;
+
+  const Symbol<Device>& device() const { return device_; }
+  int64_t storage_offset() const { return storage_offset_; }
+
   Symbol<Device>* mut_device() { return &device_; }
   void set_storage_offset(int64_t offset) { storage_offset_ = offset; }
 
-  bool operator==(const LocalTensorMeta& other) const;
+  bool operator==(const MutLocalTensorMeta& other) const;
   size_t CalcHashValue() const;
 
-  LocalTensorMeta& operator=(const LocalTensorMeta& other) = default;
+  MutLocalTensorMeta& operator=(const MutLocalTensorMeta& other) = default;
 
  private:
   Symbol<Device> device_;
@@ -127,10 +186,6 @@ class GlobalTensorMeta : public TensorMeta {
   Symbol<NdSbp> nd_sbp() const { return nd_sbp_; }
   Symbol<ParallelDesc> parallel_desc() const { return parallel_desc_; }
 
-  void set_nd_sbp(Symbol<NdSbp> val) { nd_sbp_ = val; }
-
-  void set_parallel_desc(Symbol<ParallelDesc> val) { parallel_desc_ = val; }
-
   size_t CalcHashValue() const;
 
  private:
@@ -159,4 +214,4 @@ struct hash<oneflow::one::GlobalTensorMeta> final {
 
 }  // namespace std
 
-#endif  // ONEFLOW_FRAMEWORK_TENSOR_META_H_
+#endif  // ONEFLOW_COMMON_TENSOR_META_H_

oneflow/core/eager/critical_section_phy_instr_operand.cpp

@@ -70,8 +70,7 @@ void InputCriticalSectionBeginPhyInstrOperand::AccessBlobByOpName(uint64_t of_bl
   {
     size_t header_size = of_blob->mut_blob()->blob_desc().ByteSizeOfBlobHeader();
     CHECK_EQ(header_size, eager_blob_object->shape().NumAxes() * sizeof(int64_t));
-    std::memcpy(of_blob->mut_blob()->mut_header_ptr(), eager_blob_object->mut_header_ptr(),
-                header_size);
+    CHECK_EQ(of_blob->blob().static_shape(), eager_blob_object->shape());
   }
   const auto& end_event_record = op_name2end_event_record_->at(op_name);
   if (eager_blob_object->dptr() == nullptr) {
@@ -93,7 +92,7 @@ void OutputCriticalSectionBeginPhyInstrOperand::AccessBlobByOpName(uint64_t of_b
   CHECK(interfaces_valid().at(i));
   OfBlob* of_blob = reinterpret_cast<OfBlob*>(of_blob_ptr);
   auto& eager_blob_object = eager_blob_objects_->at(i);
-  of_blob->blob().shape_view().ToShape(eager_blob_object->mut_shape());
+  CHECK_EQ(of_blob->blob().static_shape(), eager_blob_object->shape());
   const auto& end_event_record = op_name2end_event_record_->at(op_name);
   if (eager_blob_object->dptr() == nullptr) {
     end_event_record->Init(std::make_shared<NaiveEventRecord>());

oneflow/core/eager/eager_blob_object.cpp

@@ -18,32 +18,79 @@ limitations under the License.
 #include "oneflow/core/framework/to_string.h"
 #include "oneflow/core/framework/shut_down_util.h"
 #include "oneflow/core/common/shape_vec.h"
+#include "oneflow/core/common/tensor_meta.h"
 
 namespace oneflow {
+
 namespace vm {
 
-EagerBlobObject::EagerBlobObject(const std::shared_ptr<MemoryCase>& mem_case,
-                                 const std::shared_ptr<Shape>& shape,
-                                 const std::shared_ptr<Stride>& stride, DataType data_type,
-                                 const std::shared_ptr<TensorStorage>& tensor_storage,
-                                 const intrusive::shared_ptr<LocalDepObject>& dep_object)
+EagerBlobObject::EagerBlobObject(
+    const std::shared_ptr<MemoryCase>& mem_case,
+    const Symbol<one::LocalTensorMeta>& static_local_tensor_meta,
+    const std::shared_ptr<const one::MutLocalTensorMeta>& dynamic_local_tensor_meta,
+    DataType data_type, const std::shared_ptr<TensorStorage>& tensor_storage,
+    const intrusive::shared_ptr<LocalDepObject>& dep_object)
     : is_dynamic_(false),
       mem_case_(mem_case),
       data_type_(data_type),
-      shape_(shape),
-      stride_(stride),
       storage_offset_(0),
       tensor_storage_(tensor_storage),
       mem_ptr_for_allocation_compuation_pipelining_(nullptr),
       inited_mem_ptr_for_allocation_compuation_pipelining_(false),
       is_non_pod_object_placement_newed_(false),
+      pin_memory_(false),
       compute_local_dep_object_(dep_object),
-      blob_desc_(shape, stride, data_type) {
-  CHECK(static_cast<bool>(shape));
-  CHECK(static_cast<bool>(stride));
+      blob_desc_(static_cast<bool>(dynamic_local_tensor_meta)
+                     ? std::const_pointer_cast<Shape>(dynamic_local_tensor_meta->shape_ptr())
+                     : std::const_pointer_cast<Shape>(static_local_tensor_meta->shape_ptr()),
+                 static_cast<bool>(dynamic_local_tensor_meta)
+                     ? std::const_pointer_cast<Stride>(dynamic_local_tensor_meta->stride_ptr())
+                     : std::const_pointer_cast<Stride>(static_local_tensor_meta->stride_ptr()),
+                 data_type),
+      static_local_tensor_meta_(static_local_tensor_meta),
+      dynamic_local_tensor_meta_(dynamic_local_tensor_meta) {
   CHECK(static_cast<bool>(tensor_storage));
 }
 
+// user_op::TensorDesc overrides
+const Shape& EagerBlobObject::shape() const {
+  if (dynamic_local_tensor_meta_) {
+    return dynamic_local_tensor_meta_->shape();
+  } else {
+    return static_local_tensor_meta_->shape();
+  }
+}
+Shape* EagerBlobObject::mut_shape() {
+  CHECK(dynamic_local_tensor_meta_);
+  return std::const_pointer_cast<one::MutLocalTensorMeta>(dynamic_local_tensor_meta_)->mut_shape();
+}
+const Stride& EagerBlobObject::stride() const {
+  if (dynamic_local_tensor_meta_) {
+    return dynamic_local_tensor_meta_->stride();
+  } else {
+    return static_local_tensor_meta_->stride();
+  }
+}
+Stride* EagerBlobObject::mut_stride() {
+  CHECK(dynamic_local_tensor_meta_);
+  return std::const_pointer_cast<one::MutLocalTensorMeta>(dynamic_local_tensor_meta_)->mut_stride();
+}
+
+std::shared_ptr<const Shape> EagerBlobObject::shape_ptr() const {
+  if (dynamic_local_tensor_meta_) {
+    return dynamic_local_tensor_meta_->shape_ptr();
+  } else {
+    return static_local_tensor_meta_->shape_ptr();
+  }
+}
+std::shared_ptr<const Stride> EagerBlobObject::stride_ptr() const {
+  if (dynamic_local_tensor_meta_) {
+    return dynamic_local_tensor_meta_->stride_ptr();
+  } else {
+    return static_local_tensor_meta_->stride_ptr();
+  }
+}
+
 Blob* EagerBlobObject::blob() {
   if (!blob_) {
     blob_.reset(new Blob(*mem_case_, &blob_desc_, mut_header_ptr(), mut_dptr<char>()));