From c91a58dc55e31e3becf39f038f2fb8db527eedf3 Mon Sep 17 00:00:00 2001
From: hesham <h.farahat@huawei.com>
Date: Thu, 16 Apr 2020 00:09:42 -0400
Subject: [PATCH] end() iterator is causing performance problem

---
 mindspore/ccsrc/dataset/core/tensor.cc | 6 ++++++
 mindspore/ccsrc/dataset/core/tensor.h  | 4 +++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/mindspore/ccsrc/dataset/core/tensor.cc b/mindspore/ccsrc/dataset/core/tensor.cc
index 8f0eae459ad..8b70fe690d4 100644
--- a/mindspore/ccsrc/dataset/core/tensor.cc
+++ b/mindspore/ccsrc/dataset/core/tensor.cc
@@ -85,6 +85,7 @@ Tensor &Tensor::operator=(Tensor &&other) noexcept {
     shape_ = other.shape();
     type_ = other.type();
     data_ = other.StartAddr();
+    data_end_ = other.data_end_;
     data_allocator_ = std::move(other.data_allocator_);
     other.Invalidate();
   }
@@ -208,11 +209,13 @@ Tensor::~Tensor() {
     if (data_allocator_ != nullptr) {
       data_allocator_->deallocate(data_);
       data_ = nullptr;
+      data_end_ = nullptr;
     } else {
       // If we didn't have an allocator, but data_ is not null then it must
       // be a stand-alone tensor that used malloc directly.
       free(data_);
       data_ = nullptr;
+      data_end_ = nullptr;
     }
   }
 }
@@ -338,8 +341,10 @@ unsigned char *Tensor::StartAddr() {
     // on the shape and type and allocate it.
     if (data_allocator_ != nullptr) {
       data_ = data_allocator_->allocate(this->SizeInBytes());
+      data_end_ = data_ + SizeInBytes();
     } else {
       data_ = static_cast<unsigned char *>(malloc(this->SizeInBytes()));
+      data_end_ = data_ + SizeInBytes();
       if (data_ == nullptr) {
         return nullptr;
       }
@@ -362,6 +367,7 @@ void Tensor::Invalidate() {
   shape_ = TensorShape::CreateUnknownRankShape();
   type_ = DataType(DataType::DE_UNKNOWN);
   data_ = nullptr;
+  data_end_ = nullptr;
   data_allocator_ = nullptr;
 }
 
diff --git a/mindspore/ccsrc/dataset/core/tensor.h b/mindspore/ccsrc/dataset/core/tensor.h
index 3409354d196..74da40c293e 100644
--- a/mindspore/ccsrc/dataset/core/tensor.h
+++ b/mindspore/ccsrc/dataset/core/tensor.h
@@ -363,7 +363,7 @@ class Tensor {
   // @return TensorIterator
   template <typename T>
   TensorIterator<T> end() {
-    return TensorIterator<T>(data_ + SizeInBytes());
+    return TensorIterator<T>(data_end_);
   }
 
  protected:
@@ -402,6 +402,8 @@ class Tensor {
   unsigned char *data_;
   // An allocator for data_
   CharAllocPtr data_allocator_;
+  // pointer to the end of the physical data
+  unsigned char *data_end_ = nullptr;
 };
 }  // namespace dataset
 }  // namespace mindspore