Rename user op tensor shape to shape view

.github/workflows/test.yml

@@ -886,7 +886,7 @@ jobs:
               body: "<details>\n <summary>Speed stats:</summary>\n\n ``` \n${{ steps.speed.outputs.stats }}\n ``` \n\n</details>".replace(/\\n/g, '\n')
             })
       - name: Module API test
-        timeout-minutes: 50
+        timeout-minutes: 60
         if: ${{ !fromJson(matrix.cache-hit) && matrix.test-type == 'module' && !fromJson(matrix.is-distributed) }}
         run: |
           docker exec -e ONEFLOW_TEST_DIR=$PWD/python/oneflow/test/modules ${{ env.TEST_CONTAINER_NAME }} bash ci/test/generic_test_multi_client.sh

oneflow/core/common/tensor_buffer.h

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/common/shape.h"
+#include "oneflow/core/common/shape_view.h"
 #include "oneflow/core/common/data_type.h"
 
 namespace oneflow {
@@ -82,6 +83,7 @@ class TensorBuffer final {
 
   bool is_allocated() const { return bool(impl_); }
   const Shape& shape() const;
+  ShapeView shape_view() const { return shape(); }
   DataType data_type() const;
   int64_t elem_cnt() const { return shape().elem_cnt(); }
   size_t nbytes() const { return elem_cnt() * GetSizeOfDataType(data_type()); }

oneflow/core/eager/op_call_instruction_type.cpp

@@ -136,14 +136,14 @@ struct OpCallInstructionUtil final {
               : nullptr,
           [compute_ctx]() -> int64_t {
             const auto cal_memory_size = [compute_ctx](const one::ArgVec& args) -> int64_t {
-              return std::accumulate(
-                  args.begin(), args.end(), static_cast<int64_t>(0),
-                  [compute_ctx](int64_t memory_size, const auto& pair) {
-                    const auto tensor =
-                        compute_ctx->Tensor4ArgNameAndIndex(pair.first, pair.second);
-                    return memory_size
-                           + tensor->shape().elem_cnt() * GetSizeOfDataType(tensor->data_type());
-                  });
+              return std::accumulate(args.begin(), args.end(), static_cast<int64_t>(0),
+                                     [compute_ctx](int64_t memory_size, const auto& pair) {
+                                       const auto tensor = compute_ctx->Tensor4ArgNameAndIndex(
+                                           pair.first, pair.second);
+                                       return memory_size
+                                              + tensor->shape_view().elem_cnt()
+                                                    * GetSizeOfDataType(tensor->data_type());
+                                     });
             };
             return cal_memory_size(compute_ctx->inputs()) + cal_memory_size(compute_ctx->outputs());
           },

oneflow/core/framework/user_op_hob.h

@@ -46,6 +46,7 @@ ALWAYS_INLINE inline auto HobDataType(const std::string& tensor_name, int tensor
   return hob::make_custom(
       string_stream.str(), [tensor_name, tensor_idx](const KernelRegContext& ctx) -> DataType {
         const user_op::TensorDesc* desc = ctx.TensorDesc4ArgNameAndIndex(tensor_name, tensor_idx);
+        CHECK(desc != nullptr) << "key `" << tensor_name << "_" << tensor_idx << "` not found.";
         return desc->data_type();
       });
 }

oneflow/core/framework/user_op_tensor.h

@@ -38,8 +38,8 @@ class Tensor {
   ~Tensor() = default;
 #pragma GCC diagnostic pop
 
-  virtual ShapeView shape() const = 0;
-  virtual MutShapeView mut_shape() = 0;
+  virtual ShapeView shape_view() const = 0;
+  virtual MutShapeView mut_shape_view() = 0;
   virtual const Stride& stride() const = 0;
   virtual DataType data_type() const = 0;
   virtual const MemoryCase& mem_case() const = 0;

oneflow/core/kernel/blob_tensor_view.cpp

@@ -22,9 +22,9 @@ namespace user_op {
 
 BlobTensorView::BlobTensorView(Blob* blob) : blob_(blob) {}
 
-ShapeView BlobTensorView::shape() const { return blob_->shape(); }
+ShapeView BlobTensorView::shape_view() const { return blob_->shape(); }
 
-MutShapeView BlobTensorView::mut_shape() { return *blob_->mut_shape_view(); }
+MutShapeView BlobTensorView::mut_shape_view() { return *blob_->mut_shape_view(); }
 
 const Stride& BlobTensorView::stride() const { return blob_->stride(); }
 

oneflow/core/kernel/blob_tensor_view.h

@@ -29,8 +29,8 @@ class BlobTensorView final : public Tensor {
   explicit BlobTensorView(Blob* blob);
   ~BlobTensorView() = default;
 
-  ShapeView shape() const override;
-  MutShapeView mut_shape() override;
+  ShapeView shape_view() const override;
+  MutShapeView mut_shape_view() override;
   const Stride& stride() const override;
   DataType data_type() const override;
   const MemoryCase& mem_case() const override;

oneflow/core/kernel/user_kernel.cpp

@@ -427,14 +427,14 @@ class UserKernelInferContext final : public user_op::KernelInferContext {
     user_op::Tensor* arg_tensor = Tensor4ArgNameAndIndex(arg_name, arg_index);
     CHECK(arg_tensor != nullptr) << "Tensor of arg (" << arg_name << "," << arg_index
                                  << ") is not found";
-    return arg_tensor->shape();
+    return arg_tensor->shape_view();
   }
   MutShapeView MutShapeView4ArgNameAndIndex(const std::string& arg_name,
                                             int32_t arg_index) override {
     user_op::Tensor* arg_tensor = Tensor4ArgNameAndIndex(arg_name, arg_index);
     CHECK(arg_tensor != nullptr) << "Tensor of arg (" << arg_name << "," << arg_index
                                  << ") is not found";
-    return arg_tensor->mut_shape();
+    return arg_tensor->mut_shape_view();
   }
 
   user_op::InferContext* MutOpInferContext() override { return &op_infer_ctx_; }

oneflow/extension/python/py_compute.cpp

@@ -58,9 +58,9 @@ void TensorToNumpy(const user_op::Tensor* tensor, PyObject** arg_ptr) {
   int type_num = CHECK_JUST(numpy::OFDataTypeToNumpyType(tensor->data_type()));
   VLOG(3) << "Tensor data type " << DataType_Name(tensor->data_type()) << " Numpy type "
           << type_num;
-  int dim_size = tensor->shape().NumAxes();
+  int dim_size = tensor->shape_view().NumAxes();
   npy_intp dims[dim_size];
-  FOR_RANGE(size_t, i, 0, dim_size) { dims[i] = tensor->shape().At(i); }
+  FOR_RANGE(size_t, i, 0, dim_size) { dims[i] = tensor->shape_view().At(i); }
 
   void* data = TensorToMem(tensor);
   auto* np_array =
@@ -105,9 +105,9 @@ void NumpyToTensor(PyObject* arg, user_op::Tensor* tensor) {
 
   int64_t array_elem_cnt = 1;
   FOR_RANGE(int, i, 0, PyArray_NDIM(array)) { array_elem_cnt *= PyArray_SHAPE(array)[i]; }
-  CHECK_EQ(array_elem_cnt, tensor->shape().elem_cnt())
+  CHECK_EQ(array_elem_cnt, tensor->shape_view().elem_cnt())
       << "Numpy array element count " << array_elem_cnt
-      << " is not equal to OneFlow tensor element count " << tensor->shape().elem_cnt();
+      << " is not equal to OneFlow tensor element count " << tensor->shape_view().elem_cnt();
 
   void* array_data_ptr = PyArray_DATA(array);
   MemToTensor(array_data_ptr, array_elem_cnt, tensor);

oneflow/ir/oneflow-extension/extension.cpp

@@ -77,8 +77,8 @@ OpaqueMemRefDescriptor CreateMemRefDescriptor(user_op::Tensor* tensor) {
   auto desc = new MemRefType();
   *desc = mlir::detail::makeStridedMemRefDescriptor<N>(
       tensor->dptr<T>(), tensor->dptr<T>(),
-      {tensor->shape().ptr(), tensor->shape().ptr() + tensor->shape().NumAxes()},
-      {tensor->shape().ptr(), tensor->shape().ptr() + tensor->shape().NumAxes()});
+      {tensor->shape_view().ptr(), tensor->shape_view().ptr() + tensor->shape_view().NumAxes()},
+      {tensor->shape_view().ptr(), tensor->shape_view().ptr() + tensor->shape_view().NumAxes()});
   auto deleter = [](void const* data) {
     auto p = static_cast<MemRefType const*>(data);
     delete p;
@@ -92,8 +92,8 @@ OpaqueMemRefDescriptor CreateMutMemRefDescriptor(user_op::Tensor* tensor) {
   auto desc = new MemRefType();
   *desc = mlir::detail::makeStridedMemRefDescriptor<N>(
       tensor->mut_dptr<T>(), tensor->mut_dptr<T>(),
-      {tensor->shape().ptr(), tensor->shape().ptr() + tensor->shape().NumAxes()},
-      {tensor->shape().ptr(), tensor->shape().ptr() + tensor->shape().NumAxes()});
+      {tensor->shape_view().ptr(), tensor->shape_view().ptr() + tensor->shape_view().NumAxes()},
+      {tensor->shape_view().ptr(), tensor->shape_view().ptr() + tensor->shape_view().NumAxes()});
   auto deleter = [](void const* data) {
     auto p = static_cast<MemRefType const*>(data);
     delete p;
@@ -120,13 +120,13 @@ llvm::SmallVector<OpaqueMemRefDescriptor> GetMLIRCInterfaceArgs(
   for (auto& pair : ctx->inputs()) {
     auto tensor = ctx->Tensor4ArgNameAndIndex(pair.first, pair.second);
     auto ref = SwitchCreateMemRefDescriptor(
-        SwitchCase(tensor->shape().NumAxes(), tensor->data_type()), tensor);
+        SwitchCase(tensor->shape_view().NumAxes(), tensor->data_type()), tensor);
     args.push_back(ref);
   }
   for (auto& pair : ctx->outputs()) {
     auto tensor = ctx->Tensor4ArgNameAndIndex(pair.first, pair.second);
     auto ref = SwitchCreateMutMemRefDescriptor(
-        SwitchCase(tensor->shape().NumAxes(), tensor->data_type()), tensor);
+        SwitchCase(tensor->shape_view().NumAxes(), tensor->data_type()), tensor);
     args.push_back(ref);
   }
   return args;

oneflow/user/data/coco_parser.cpp

@@ -64,30 +64,31 @@ void COCOParser::Parse(BatchType& batch_data, user_op::KernelComputeContext* ctx
     }
   });
   // dynamic batch size
-  if (image_tensor->shape().elem_cnt() != batch_data.size()) {
-    CHECK_EQ(image_tensor->shape().NumAxes(), 1);
-    image_tensor->mut_shape().Set(0, batch_data.size());
+  if (image_tensor->shape_view().elem_cnt() != batch_data.size()) {
+    CHECK_EQ(image_tensor->shape_view().NumAxes(), 1);
+    image_tensor->mut_shape_view().Set(0, batch_data.size());
   }
-  if (image_id_tensor && image_id_tensor->shape().At(0) != batch_data.size()) {
-    image_id_tensor->mut_shape().Set(0, batch_data.size());
+  if (image_id_tensor && image_id_tensor->shape_view().At(0) != batch_data.size()) {
+    image_id_tensor->mut_shape_view().Set(0, batch_data.size());
   }
-  if (image_size_tensor && image_size_tensor->shape().At(0) != batch_data.size()) {
-    image_size_tensor->mut_shape().Set(0, batch_data.size());
+  if (image_size_tensor && image_size_tensor->shape_view().At(0) != batch_data.size()) {
+    image_size_tensor->mut_shape_view().Set(0, batch_data.size());
   }
-  if (bbox_tensor && bbox_tensor->shape().elem_cnt() != batch_data.size()) {
-    CHECK_EQ(bbox_tensor->shape().NumAxes(), 1);
-    bbox_tensor->mut_shape().Set(0, batch_data.size());
+  if (bbox_tensor && bbox_tensor->shape_view().elem_cnt() != batch_data.size()) {
+    CHECK_EQ(bbox_tensor->shape_view().NumAxes(), 1);
+    bbox_tensor->mut_shape_view().Set(0, batch_data.size());
   }
-  if (label_tensor && label_tensor->shape().elem_cnt() != batch_data.size()) {
-    CHECK_EQ(label_tensor->shape().NumAxes(), 1);
-    label_tensor->mut_shape().Set(0, batch_data.size());
+  if (label_tensor && label_tensor->shape_view().elem_cnt() != batch_data.size()) {
+    CHECK_EQ(label_tensor->shape_view().NumAxes(), 1);
+    label_tensor->mut_shape_view().Set(0, batch_data.size());
   }
-  if (segm_tensor && segm_index_tensor && segm_tensor->shape().elem_cnt() != batch_data.size()) {
-    CHECK_EQ(segm_tensor->shape().NumAxes(), 1);
-    CHECK_EQ(segm_index_tensor->shape().NumAxes(), 1);
-    CHECK_EQ(segm_tensor->shape().elem_cnt(), segm_index_tensor->shape().elem_cnt());
-    segm_tensor->mut_shape().Set(0, batch_data.size());
-    segm_index_tensor->mut_shape().Set(0, batch_data.size());
+  if (segm_tensor && segm_index_tensor
+      && segm_tensor->shape_view().elem_cnt() != batch_data.size()) {
+    CHECK_EQ(segm_tensor->shape_view().NumAxes(), 1);
+    CHECK_EQ(segm_index_tensor->shape_view().NumAxes(), 1);
+    CHECK_EQ(segm_tensor->shape_view().elem_cnt(), segm_index_tensor->shape_view().elem_cnt());
+    segm_tensor->mut_shape_view().Set(0, batch_data.size());
+    segm_index_tensor->mut_shape_view().Set(0, batch_data.size());
   }
 }
 

oneflow/user/data/ofrecord_image_classification_dataset.cpp

@@ -103,7 +103,7 @@ void DecodeWorker(const std::string& image_feature_name, const std::string& labe
     CHECK(receive_status == kBufferStatusSuccess);
     OFRecord record;
     CHECK(record.ParseFromArray(serialized_record.data<char>(),
-                                serialized_record.shape().elem_cnt()));
+                                serialized_record.shape_view().elem_cnt()));
     ImageClassificationDataInstance instance;
     DecodeImageFromOFRecord(record, image_feature_name, color_space, &instance.image);
     DecodeLabelFromFromOFRecord(record, label_feature_name, &instance.label);

oneflow/user/data/ofrecord_image_classification_parser.h

@@ -38,12 +38,12 @@ class OFRecordImageClassificationParser final : public Parser<ImageClassificatio
   void Parse(BatchType& batch_data, user_op::KernelComputeContext* ctx) override {
     const int64_t batch_size = batch_data.size();
     user_op::Tensor* image_tensor = ctx->Tensor4ArgNameAndIndex("image", 0);
-    CHECK_EQ(image_tensor->shape().NumAxes(), 1);
-    CHECK_EQ(image_tensor->shape().At(0), batch_size);
+    CHECK_EQ(image_tensor->shape_view().NumAxes(), 1);
+    CHECK_EQ(image_tensor->shape_view().At(0), batch_size);
     auto* image_buffers = image_tensor->mut_dptr<TensorBuffer>();
     user_op::Tensor* label_tensor = ctx->Tensor4ArgNameAndIndex("label", 0);
-    CHECK_EQ(label_tensor->shape().NumAxes(), 1);
-    CHECK_EQ(label_tensor->shape().At(0), batch_size);
+    CHECK_EQ(label_tensor->shape_view().NumAxes(), 1);
+    CHECK_EQ(label_tensor->shape_view().At(0), batch_size);
     auto* label_buffers = label_tensor->mut_dptr<TensorBuffer>();
     for (size_t i = 0; i < batch_data.size(); ++i) {
       auto& instance = batch_data[i];

oneflow/user/data/ofrecord_parser.h

@@ -40,9 +40,9 @@ class OFRecordParser final : public Parser<TensorBuffer> {
       auto& sample = batch_data[i];
       CHECK(dptr[i].ParseFromArray(sample.data(), sample.nbytes()));
     });
-    if (batch_data.size() != out_tensor->shape().elem_cnt()) {
-      CHECK_EQ(out_tensor->mut_shape().NumAxes(), 1);
-      out_tensor->mut_shape().Set(0, batch_data.size());
+    if (batch_data.size() != out_tensor->shape_view().elem_cnt()) {
+      CHECK_EQ(out_tensor->mut_shape_view().NumAxes(), 1);
+      out_tensor->mut_shape_view().Set(0, batch_data.size());
     }
   }
 };

oneflow/user/image/image_util.cpp

@@ -39,10 +39,10 @@ void ImageUtil::ConvertColor(const std::string& input_color, const cv::Mat& inpu
 }
 
 cv::Mat GenCvMat4ImageBuffer(const TensorBuffer& image_buffer) {
-  CHECK_EQ(image_buffer.shape().NumAxes(), 3);
-  int h = image_buffer.shape().At(0);
-  int w = image_buffer.shape().At(1);
-  int channels = image_buffer.shape().At(2);
+  CHECK_EQ(image_buffer.shape_view().NumAxes(), 3);
+  int h = image_buffer.shape_view().At(0);
+  int w = image_buffer.shape_view().At(1);
+  int channels = image_buffer.shape_view().At(2);
   DataType data_type = image_buffer.data_type();
   if (channels == 1 && data_type == DataType::kUInt8) {
     return CreateMatWithPtr(h, w, CV_8UC1, image_buffer.data<uint8_t>());
@@ -60,19 +60,19 @@ cv::Mat GenCvMat4ImageBuffer(const TensorBuffer& image_buffer) {
 
 cv::Mat GenCvMat4ImageTensor(const user_op::Tensor* image_tensor, int image_offset) {
   int has_batch_dim = 0;
-  if (image_tensor->shape().NumAxes() == 3) {
+  if (image_tensor->shape_view().NumAxes() == 3) {
     has_batch_dim = 0;
     image_offset = 0;
-  } else if (image_tensor->shape().NumAxes() == 4) {
+  } else if (image_tensor->shape_view().NumAxes() == 4) {
     has_batch_dim = 1;
     CHECK_GE(image_offset, 0);
-    CHECK_LT(image_offset, image_tensor->shape().At(0));
+    CHECK_LT(image_offset, image_tensor->shape_view().At(0));
   } else {
     UNIMPLEMENTED();
   }
-  int h = image_tensor->shape().At(0 + has_batch_dim);
-  int w = image_tensor->shape().At(1 + has_batch_dim);
-  int c = image_tensor->shape().At(2 + has_batch_dim);
+  int h = image_tensor->shape_view().At(0 + has_batch_dim);
+  int w = image_tensor->shape_view().At(1 + has_batch_dim);
+  int c = image_tensor->shape_view().At(2 + has_batch_dim);
   int elem_offset = image_offset * h * w * c;
   DataType data_type = image_tensor->data_type();
   if (c == 1 && data_type == DataType::kUInt8) {

oneflow/user/kernels/acc_kernel.cpp

@@ -31,13 +31,13 @@ class AccKernel final : public user_op::OpKernel {
   void Compute(user_op::KernelComputeContext* ctx) const override {
     const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0);
     user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
-    CHECK_EQ(in->shape().elem_cnt(), out->shape().elem_cnt());
+    CHECK_EQ(in->shape_view().elem_cnt(), out->shape_view().elem_cnt());
     CHECK_EQ(in->data_type(), out->data_type());
     std::unique_ptr<ep::primitive::Add> primitive =
         ep::primitive::NewPrimitive<ep::primitive::AddFactory>(ctx->device_type(), in->data_type());
     CHECK(primitive);
     primitive->Launch(ctx->stream(), out->dptr(), in->dptr(), out->mut_dptr(),
-                      in->shape().elem_cnt());
+                      in->shape_view().elem_cnt());
   }
   bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
 };

oneflow/user/kernels/adaptive_pool_cpu_kernel.cpp

@@ -108,7 +108,7 @@ void AvgBackwardCompute(user_op::KernelComputeContext* ctx, const int32_t& dim)
   const T* out_ptr = grad_output->dptr<T>();
   T* in_ptr = grad_input->mut_dptr<T>();
 
-  std::fill(in_ptr, in_ptr + grad_input->shape().elem_cnt(), static_cast<T>(0));
+  std::fill(in_ptr, in_ptr + grad_input->shape_view().elem_cnt(), static_cast<T>(0));
 
   const int64_t input_width = in.Count(4);
   const int64_t output_width = out.Count(4);

oneflow/user/kernels/adaptive_pool_gpu_kernel.cu

@@ -150,7 +150,7 @@ void AvgForwardCompute(KernelComputeContext* ctx, const int32_t& dim) {
   const Shape& in = GetShape5D(x_shape, data_format, dim);
   const Shape& out = GetShape5D(y_shape, data_format, dim);
 
-  const int out_elems = out_tensor->shape().elem_cnt();
+  const int out_elems = out_tensor->shape_view().elem_cnt();
 
   RUN_CUDA_KERNEL((AdaptiveAvgPoolCudaKernel<T>), ctx->stream(), out_elems, in_ptr, out_ptr,
                   out_elems, in.At(2), in.At(3), in.At(4), out.At(2), out.At(3), out.At(4));
@@ -171,8 +171,8 @@ void AvgBackwardCompute(KernelComputeContext* ctx, const int32_t& dim) {
   const Shape& in = GetShape5D(dx_shape, data_format, dim);
   const Shape& out = GetShape5D(dy_shape, data_format, dim);
 
-  const int in_elems = in_tensor->shape().elem_cnt();
-  const int out_elems = out_tensor->shape().elem_cnt();
+  const int in_elems = in_tensor->shape_view().elem_cnt();
+  const int out_elems = out_tensor->shape_view().elem_cnt();
 
   RUN_CUDA_KERNEL((InitPtr<T>), ctx->stream(), in_elems, in_elems, in_ptr);
   RUN_CUDA_KERNEL((AdaptiveAvgPoolGradCudaKernel<T>), ctx->stream(), out_elems, in_ptr, out_ptr,

oneflow/user/kernels/add_n_kernel.cpp

@@ -44,13 +44,13 @@ class AddNKernel : public OpKernel, public CudaGraphSupport {
     CHECK(primitive);
     Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
     const DataType data_type = out->data_type();
-    const size_t count = out->shape().elem_cnt();
+    const size_t count = out->shape_view().elem_cnt();
     if (count == 0) { return; }
     size_t in_num = ctx->inputs().size();
     std::vector<const void*> srcs(in_num);
     for (size_t i = 0; i < in_num; ++i) {
       const Tensor* in_i = ctx->Tensor4ArgNameAndIndex("in", i);
-      CHECK_EQ(in_i->shape().elem_cnt(), count);
+      CHECK_EQ(in_i->shape_view().elem_cnt(), count);
       CHECK_EQ(in_i->data_type(), data_type);
       srcs[i] = in_i->template dptr();
     }

oneflow/user/kernels/affine_grid_kernel.cpp

@@ -38,9 +38,9 @@ class AffineGridKernel final : public user_op::OpKernel {
     bool is_2d_grid = true;
     if (size.NumAxes() == 5) { is_2d_grid = false; }
 
-    int64_t N = theta->shape().At(0);
-    int64_t theta_h = theta->shape().At(1);
-    int64_t theta_w = theta->shape().At(2);
+    int64_t N = theta->shape_view().At(0);
+    int64_t theta_h = theta->shape_view().At(1);
+    int64_t theta_w = theta->shape_view().At(2);
 
     if (is_2d_grid) {
       int64_t H = size.At(2);
@@ -108,9 +108,9 @@ class AffineGridGradKernel final : public user_op::OpKernel {
     bool is_2d_grid = true;
     if (size.NumAxes() == 5) { is_2d_grid = false; }
 
-    int64_t N = dtheta->shape().At(0);
-    int64_t dtheta_h = dtheta->shape().At(1);
-    int64_t dtheta_w = dtheta->shape().At(2);
+    int64_t N = dtheta->shape_view().At(0);
+    int64_t dtheta_h = dtheta->shape_view().At(1);
+    int64_t dtheta_w = dtheta->shape_view().At(2);
 
     if (is_2d_grid) {
       int64_t H = size.At(2);