【Hackathon 7th No.25】为 Paddle 新增 sparse_dim/dense_dim -part (#69132)

NKNaN · web-flow · commit 121745f1f724 · 2024-11-28T17:22:26.000+08:00
* add sparse dim dense dim

* fix example code

* fix test

* resolve conflict

* update pir
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -3316,6 +3316,121 @@ static PyObject* tensor_is_contiguous(TensorObject* self,
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+PyDoc_STRVAR(tensor_method_sparse_dim__doc__,
+             R"DOC(sparse_dim($self, /)
+--
+
+Returns the number of sparse dimensions of sparse Tensor.
+
+Note:
+    **If self is not sparse Tensor, return 0.**
+
+Returns:
+    int, sparse dim of self Tensor
+
+Examples:
+
+    .. code-block:: python
+
+        >>> import paddle
+
+        >>> indices = [[0, 1, 2], [1, 2, 0]]
+        >>> values = [1.0, 2.0, 3.0]
+        >>> dense_shape = [3, 3]
+        >>> coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape)
+        >>> coo.sparse_dim()
+        2
+
+        >>> crows = [0, 2, 3, 5]
+        >>> cols = [1, 3, 2, 0, 1]
+        >>> values = [1, 2, 3, 4, 5]
+        >>> dense_shape = [3, 4]
+        >>> csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape)
+        >>> csr.sparse_dim()
+        2
+
+        >>> dense = paddle.to_tensor([1, 2, 3])
+        >>> dense.sparse_dim()
+        0
+
+)DOC");  // NOLINT
+
+static PyObject* tensor_method_sparse_dim(TensorObject* self,
+                                          PyObject* args,
+                                          PyObject* kwargs) {
+  EAGER_TRY
+  if (self->tensor.is_sparse_coo_tensor()) {
+    auto sparse_coo_tensor =
+        std::dynamic_pointer_cast<phi::SparseCooTensor>(self->tensor.impl());
+    return ToPyObject(sparse_coo_tensor->sparse_dim());
+  } else if (self->tensor.is_sparse_csr_tensor()) {
+    auto sparse_csr_tensor =
+        std::dynamic_pointer_cast<phi::SparseCsrTensor>(self->tensor.impl());
+    return ToPyObject(sparse_csr_tensor->sparse_dim());
+  } else {
+    return ToPyObject(0);
+  }
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
+PyDoc_STRVAR(tensor_method_dense_dim__doc__,
+             R"DOC(dense_dim($self, /)
+--
+
+Returns the number of dense dimensions of sparse Tensor.
+
+Note:
+    **If self is not sparse Tensor, return len(self.shape).**
+
+Returns:
+    int, dense dim of self Tensor
+
+Examples:
+
+    .. code-block:: python
+
+        >>> import paddle
+        >>> import numpy as np
+
+        >>> indices = [[0, 1, 1], [2, 0, 2]]
+        >>> values = np.array([[3, 4], [5, 6], [7, 8]])
+        >>> dense_shape = [2, 3, 2]
+        >>> coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape)
+        >>> coo.dense_dim()
+        1
+
+        >>> crows = [0, 2, 3, 5]
+        >>> cols = [1, 3, 2, 0, 1]
+        >>> values = [1, 2, 3, 4, 5]
+        >>> dense_shape = [3, 4]
+        >>> csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape)
+        >>> csr.dense_dim()
+        0
+
+        >>> dense = paddle.to_tensor([[1, 2, 3]])
+        >>> dense.dense_dim()
+        >>> 2
+
+)DOC");  // NOLINT
+
+static PyObject* tensor_method_dense_dim(TensorObject* self,
+                                         PyObject* args,
+                                         PyObject* kwargs) {
+  EAGER_TRY
+  if (self->tensor.is_sparse_coo_tensor()) {
+    auto sparse_coo_tensor =
+        std::dynamic_pointer_cast<phi::SparseCooTensor>(self->tensor.impl());
+    return ToPyObject(sparse_coo_tensor->dense_dim());
+  } else if (self->tensor.is_sparse_csr_tensor()) {
+    auto sparse_csr_tensor =
+        std::dynamic_pointer_cast<phi::SparseCsrTensor>(self->tensor.impl());
+    return ToPyObject(sparse_csr_tensor->dense_dim());
+  } else {
+    return ToPyObject(self->tensor.shape().size());
+  }
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* tensor_method__set_impl(TensorObject* self,
                                          PyObject* args,
                                          PyObject* kwargs) {
@@ -3581,6 +3696,14 @@ PyMethodDef variable_methods[] = {  // NOLINT
      (PyCFunction)(void (*)())tensor_method_is_coalesced,
      METH_VARARGS | METH_KEYWORDS,
      tensor_is_coalesced__doc__},
+    {"sparse_dim",
+     (PyCFunction)(void (*)())tensor_method_sparse_dim,
+     METH_VARARGS | METH_KEYWORDS,
+     tensor_method_sparse_dim__doc__},
+    {"dense_dim",
+     (PyCFunction)(void (*)())tensor_method_dense_dim,
+     METH_VARARGS | METH_KEYWORDS,
+     tensor_method_dense_dim__doc__},
     /***the method of sparse tensor****/
     {"element_size",
      (PyCFunction)(void (*)())tensor_method_element_size,
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1461,9 +1461,42 @@ void BindValue(py::module *m) {
               return py::cast<py::none>(Py_None);
             }
           })
-      .def("_clone", [](Value self) {
-        // Return a new value owned by python side
-        return self;
+      .def("_clone",
+           [](Value self) {
+             // Return a new value owned by python side
+             return self;
+           })
+      .def("sparse_dim",
+           [](Value self) -> int32_t {
+             auto op_result = self.dyn_cast<OpResult>();
+             pir::Operation *operation = op_result.owner();
+             if (self.type().isa<SparseCooTensorType>() &&
+                 operation->name() == "pd_op.sparse_coo_tensor_sp") {
+               std::vector<Value> sources = operation->operands_source();
+               Value non_zero_indices = sources[1];
+               return phi::vectorize(GetValueDims(non_zero_indices))[0];
+             } else if (self.type().isa<SparseCsrTensorType>()) {
+               PADDLE_THROW(common::errors::InvalidType(
+                   "SparseCsrTensor is unsupported in pir mode."));
+             } else {
+               return 0;
+             }
+           })
+      .def("dense_dim", [](Value self) -> int32_t {
+        auto op_result = self.dyn_cast<OpResult>();
+        pir::Operation *operation = op_result.owner();
+        if (self.type().isa<SparseCooTensorType>() &&
+            operation->name() == "pd_op.sparse_coo_tensor_sp") {
+          std::vector<Value> sources = operation->operands_source();
+          Value non_zero_indices = sources[1];
+          int32_t dims = phi::vectorize(GetValueDims(self)).size();
+          return dims - phi::vectorize(GetValueDims(non_zero_indices))[0];
+        } else if (self.type().isa<SparseCsrTensorType>()) {
+          PADDLE_THROW(common::errors::InvalidType(
+              "SparseCsrTensor is unsupported in pir mode."));
+        } else {
+          return phi::vectorize(GetValueDims(self)).size();
+        }
       });
 }
 
diff --git a/paddle/phi/core/sparse_csr_tensor.cc b/paddle/phi/core/sparse_csr_tensor.cc
@@ -160,4 +160,16 @@ void SparseCsrTensor::set_meta(const SparseTensorMeta& meta) {
   meta_.dtype = meta.dtype;
   meta_.layout = meta.layout;
 }
+
+int32_t SparseCsrTensor::sparse_dim() const { return 2; }
+
+int32_t SparseCsrTensor::dense_dim() const {
+  int32_t nze_dim = this->non_zero_elements_.dims().size();
+  int32_t batch_dim = this->non_zero_crows_.dims().size() - 1;
+  // layout of SparseCsrTensor has not been implemented yet
+  // int32_t block_dim =  = (layout_ == kSparseBsr || layout_ == kSparseBsc ? 2
+  // : 0);
+  int32_t block_dim = 0;
+  return nze_dim - batch_dim - block_dim - 1;
+}
 }  // namespace phi
diff --git a/paddle/phi/core/sparse_csr_tensor.h b/paddle/phi/core/sparse_csr_tensor.h
@@ -110,6 +110,14 @@ class SparseCsrTensor : public TensorBase,
   /// \return The data type of the tensor.
   DataType dtype() const noexcept override { return meta_.dtype; }
 
+  /// \brief get the sparse dim
+  /// \return The sparse dim of the tensor.
+  int32_t sparse_dim() const;
+
+  /// \brief get the dense dim
+  /// \return The dense dim of the tensor.
+  int32_t dense_dim() const;
+
 #ifndef PADDLE_WITH_CUSTOM_KERNEL
   void set_type(const DataType dtype);
 #endif
diff --git a/python/paddle/tensor/tensor.prototype.pyi b/python/paddle/tensor/tensor.prototype.pyi
@@ -225,6 +225,7 @@ class AbstractTensor:
     @data.setter
     def data(self, value: Tensor) -> None: ...
     def data_ptr(self) -> int: ...
+    def dense_dim(self) -> int: ...
     def detach(self) -> Tensor: ...
     def detach_(self) -> Tensor: ...
     @property
@@ -286,6 +287,7 @@ class AbstractTensor:
     def shape(self) -> list[int]: ...
     @property
     def size(self) -> int: ...
+    def sparse_dim(self) -> int: ...
     @property
     def stop_gradient(self) -> bool: ...
     @stop_gradient.setter
diff --git a/test/legacy_test/test_dense_dim.py b/test/legacy_test/test_dense_dim.py
diff --git a/test/legacy_test/test_sparse_dim.py b/test/legacy_test/test_sparse_dim.py
