Support SparseCooTensorType (#62868)

risemeup1 · web-flow · commit cc53f1cd7f6a · 2024-03-21T00:40:50.000+08:00
* support sparsecootensortype

* support sparsecootensortype

* support sparsecootensortype

* support sparsecootensortype

* support sparsecootensortype

* support sparsecootensortype

* support sparsecootensortype
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc
@@ -299,6 +299,7 @@ void PrintOperationImpl(pir::Operation* op,
 
 void OperatorDialect::initialize() {
   RegisterTypes<paddle::dialect::SelectedRowsType,
+                paddle::dialect::SparseCooTensorType,
                 paddle::dialect::DenseTensorArrayType>();
 
   RegisterAttributes<paddle::dialect::IntArrayAttribute,
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_type.cc b/paddle/fluid/pir/dialect/operator/ir/op_type.cc
@@ -78,8 +78,51 @@ DenseTensorArrayType DenseTensorArrayType::dyn_cast_impl(Type type) {
   return nullptr;
 }
 
+pir::Type SparseCooTensorType::dtype() const { return storage()->dtype_; }
+
+const common::DDim& SparseCooTensorType::dims() const {
+  return storage()->dims_;
+}
+
+const common::DDim& SparseCooTensorType::non_zero_dims() const {
+  return storage()->non_zero_dims_;
+}
+
+common::DataLayout SparseCooTensorType::data_layout() const {
+  return storage()->layout_;
+}
+
+pir::DenseTensorType SparseCooTensorType::non_zero_indices() const {
+  return storage()->non_zero_indices_;
+}
+
+pir::DenseTensorType SparseCooTensorType::non_zero_elements() const {
+  return storage()->non_zero_elements_;
+}
+
+bool SparseCooTensorType::coalesced() const { return storage()->coalesced_; }
+
+bool SparseCooTensorType::classof(Type type) {
+  if (type) {
+    if (type.type_id() == type_id()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+SparseCooTensorType SparseCooTensorType::dyn_cast_impl(Type type) {
+  if (type) {
+    if (type.type_id() == type_id()) {
+      return SparseCooTensorType(type.storage());
+    }
+  }
+  return nullptr;
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SelectedRowsType)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::DenseTensorArrayType)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SparseCooTensorType)
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_type.h b/paddle/fluid/pir/dialect/operator/ir/op_type.h
@@ -74,8 +74,50 @@ class DenseTensorArrayType
   static DenseTensorArrayType dyn_cast_impl(Type type);
 };
 
+class IR_API SparseCooTensorType
+    : public pir::Type::
+          TypeBase<SparseCooTensorType, pir::Type, SparseCooTensorTypeStorage> {
+ public:
+  using Base::Base;
+
+  pir::Type dtype() const;
+  const common::DDim &dims() const;
+  const common::DDim &non_zero_dims() const;
+  common::DataLayout data_layout() const;
+  pir::DenseTensorType non_zero_indices() const;
+  pir::DenseTensorType non_zero_elements() const;
+  bool coalesced() const;
+
+  ///
+  /// \brief Implementation of 'classof' that compares the type id of
+  /// the provided value with the concrete type id.
+  ///
+  static bool classof(pir::Type type);
+
+  static SparseCooTensorType dyn_cast_impl(pir::Type type);
+
+  static SparseCooTensorType get(pir::IrContext *ctx,
+                                 pir::Type dtype,
+                                 const common::DDim &dims,
+                                 const common::DDim &non_zero_dims,
+                                 common::DataLayout layout,
+                                 pir::DenseTensorType non_zero_indices,
+                                 pir::DenseTensorType non_zero_elements,
+                                 bool coalesced = false) {
+    return Base::get(ctx,
+                     dtype,
+                     dims,
+                     non_zero_dims,
+                     layout,
+                     non_zero_indices,
+                     non_zero_elements,
+                     coalesced);
+  }
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SelectedRowsType)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::DenseTensorArrayType)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SparseCooTensorType)
diff --git a/paddle/fluid/pir/dialect/operator/ir/type_storage.h b/paddle/fluid/pir/dialect/operator/ir/type_storage.h
@@ -17,6 +17,7 @@
 #include <type_traits>
 
 #include "paddle/phi/core/tensor_meta.h"
+#include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/core/builtin_type_storage.h"
 #include "paddle/pir/include/core/type.h"
 #include "paddle/pir/include/core/type_base.h"
@@ -166,5 +167,124 @@ struct DenseTensorArrayTypeStorage : public pir::TypeStorage {
   phi::DataLayout layout_;
 };
 
+struct SparseCooTensorTypeStorage : public pir::TypeStorage {
+  ///
+  /// \brief Declare ParamKey according to parameter type.
+  ///
+  using ParamKey = std::tuple<pir::Type,
+                              common::DDim,
+                              common::DDim,
+                              common::DataLayout,
+                              pir::DenseTensorType,
+                              pir::DenseTensorType,
+                              bool>;
+  SparseCooTensorTypeStorage(pir::Type dtype,
+                             common::DDim dims,
+                             common::DDim non_zero_dims,
+                             common::DataLayout layout,
+                             pir::DenseTensorType non_zero_indices,
+                             pir::DenseTensorType non_zero_elements,
+                             bool coalesced = false)
+      : dtype_(dtype),
+        dims_(dims),
+        non_zero_dims_(non_zero_dims),
+        layout_(layout),
+        non_zero_indices_(non_zero_indices),
+        non_zero_elements_(non_zero_elements),
+        coalesced_(coalesced) {}
+
+  ///
+  /// \brief Each derived TypeStorage must define a Construct method, which
+  /// StorageManager uses to construct a derived TypeStorage.
+  ///
+  static SparseCooTensorTypeStorage* Construct(const ParamKey& key) {
+    return new SparseCooTensorTypeStorage(std::get<0>(key),
+                                          std::get<1>(key),
+                                          std::get<2>(key),
+                                          std::get<3>(key),
+                                          std::get<4>(key),
+                                          std::get<5>(key),
+                                          std::get<6>(key));
+  }
+
+  ///
+  /// \brief Each derived TypeStorage must provide a HashValue method.
+  ///
+  static std::size_t HashValue(const ParamKey& key) {
+    std::size_t hash_value = 0;
+    // hash dtype
+    hash_value = pir::detail::hash_combine(
+        hash_value, std::hash<pir::Type>()(std::get<0>(key)));
+    // hash dims
+    hash_value = pir::detail::hash_combine(
+        hash_value, std::hash<common::DDim>()(std::get<1>(key)));
+    // hash non_zero_dims
+    hash_value = pir::detail::hash_combine(
+        hash_value, std::hash<common::DDim>()(std::get<2>(key)));
+    // hash layout
+    hash_value = pir::detail::hash_combine(
+        hash_value,
+        std::hash<std::underlying_type<DataLayout>::type>()(
+            static_cast<std::underlying_type<DataLayout>::type>(
+                std::get<3>(key))));
+    // hash DenseTensorType
+    auto tuple1 = std::make_tuple(std::get<4>(key).dtype(),
+                                  std::get<4>(key).dims(),
+                                  std::get<4>(key).data_layout(),
+                                  std::get<4>(key).lod(),
+                                  std::get<4>(key).offset());
+    hash_value = pir::detail::hash_combine(
+        hash_value, DenseTensorTypeStorage::HashValue(tuple1));
+    // hash DenseTensorType
+    auto tuple2 = std::make_tuple(std::get<5>(key).dtype(),
+                                  std::get<5>(key).dims(),
+                                  std::get<5>(key).data_layout(),
+                                  std::get<5>(key).lod(),
+                                  std::get<5>(key).offset());
+    hash_value = pir::detail::hash_combine(
+        hash_value, DenseTensorTypeStorage::HashValue(tuple2));
+    // hash coalesced
+    hash_value = pir::detail::hash_combine(hash_value,
+                                           std::hash<bool>()(std::get<6>(key)));
+
+    return hash_value;
+  }
+
+  ///
+  /// \brief Each derived TypeStorage needs to overload operator==.
+  ///
+  bool operator==(const ParamKey& key) const {
+    return ParamKey(dtype_,
+                    dims_,
+                    non_zero_dims_,
+                    layout_,
+                    non_zero_indices_,
+                    non_zero_elements_,
+                    coalesced_) == key;
+  }
+
+  ParamKey GetAsKey() const {
+    return ParamKey(dtype_,
+                    dims_,
+                    non_zero_dims_,
+                    layout_,
+                    non_zero_indices_,
+                    non_zero_elements_,
+                    coalesced_);
+  }
+
+  ///
+  /// \brief SparseCooTensorTypeStorage include six parameters: dims, dtype,
+  /// layout, non_zero_indices_, non_zero_elements_,coalesced_.
+  ///
+
+  pir::Type dtype_;
+  common::DDim dims_;
+  common::DDim non_zero_dims_;
+  common::DataLayout layout_{DataLayout::NCHW};
+  pir::DenseTensorType non_zero_indices_;
+  pir::DenseTensorType non_zero_elements_;
+  bool coalesced_ = false;
+};
 }  // namespace dialect
 }  // namespace paddle
diff --git a/test/cpp/pir/core/type_test.cc b/test/cpp/pir/core/type_test.cc
@@ -249,6 +249,41 @@ TEST(type_test, custom_type_dialect) {
   EXPECT_EQ(dialect_integer1, dialect_integer2);
 }
 
+TEST(type_test, sparse_coo) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  pir::Type fp32_dtype = pir::Float32Type::get(ctx);
+  common::DDim dims = {4, 4};
+  common::DDim non_zero_dims = {4, 1};
+  common::DataLayout data_layout = common::DataLayout::NCHW;
+  pir::LoD lod = {{0, 1, 2}};
+  size_t offset = 0;
+  pir::DenseTensorType none_zero_indices = pir::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+  pir::DenseTensorType none_zero_elements = pir::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+  bool coalesced = false;
+  pir::Type pir_type =
+      paddle::dialect::SparseCooTensorType::get(ctx,
+                                                fp32_dtype,
+                                                dims,
+                                                non_zero_dims,
+                                                data_layout,
+                                                none_zero_indices,
+                                                none_zero_elements,
+                                                coalesced);
+
+  EXPECT_EQ(pir_type.isa<paddle::dialect::SparseCooTensorType>(), true);
+  paddle::dialect::SparseCooTensorType sparse_coo_tensor_type =
+      pir_type.dyn_cast<paddle::dialect::SparseCooTensorType>();
+  EXPECT_EQ(sparse_coo_tensor_type.dims(), dims);
+  EXPECT_EQ(sparse_coo_tensor_type.non_zero_dims(), non_zero_dims);
+  EXPECT_EQ(sparse_coo_tensor_type.data_layout(), data_layout);
+  EXPECT_EQ(sparse_coo_tensor_type.non_zero_indices(), none_zero_indices);
+  EXPECT_EQ(sparse_coo_tensor_type.non_zero_elements(), none_zero_elements);
+  EXPECT_EQ(sparse_coo_tensor_type.coalesced(), coalesced);
+}
+
 TEST(type_test, pd_op_dialect) {
   pir::IrContext *ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
