suport SparseCsrTensorType (#62894)

risemeup1 · web-flow · commit 73b45c80710e · 2024-03-21T10:25:27.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc
@@ -300,6 +300,7 @@ void PrintOperationImpl(pir::Operation* op,
 void OperatorDialect::initialize() {
   RegisterTypes<paddle::dialect::SelectedRowsType,
                 paddle::dialect::SparseCooTensorType,
+                paddle::dialect::SparseCsrTensorType,
                 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
@@ -120,9 +120,49 @@ SparseCooTensorType SparseCooTensorType::dyn_cast_impl(Type type) {
   return nullptr;
 }
 
+pir::Type SparseCsrTensorType::dtype() const { return storage()->dtype_; }
+
+const common::DDim& SparseCsrTensorType::dims() const {
+  return storage()->dims_;
+}
+
+common::DataLayout SparseCsrTensorType::data_layout() const {
+  return storage()->layout_;
+}
+
+pir::DenseTensorType SparseCsrTensorType::non_zero_crows() const {
+  return storage()->non_zero_crows_;
+}
+
+pir::DenseTensorType SparseCsrTensorType::non_zero_cols() const {
+  return storage()->non_zero_cols_;
+}
+
+pir::DenseTensorType SparseCsrTensorType::non_zero_elements() const {
+  return storage()->non_zero_elements_;
+}
+
+bool SparseCsrTensorType::classof(Type type) {
+  if (type) {
+    if (type.type_id() == type_id()) {
+      return true;
+    }
+  }
+  return false;
+}
+
+SparseCsrTensorType SparseCsrTensorType::dyn_cast_impl(Type type) {
+  if (type) {
+    if (type.type_id() == type_id()) {
+      return SparseCsrTensorType(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)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SparseCsrTensorType)
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_type.h b/paddle/fluid/pir/dialect/operator/ir/op_type.h
@@ -115,9 +115,48 @@ class IR_API SparseCooTensorType
   }
 };
 
+class IR_API SparseCsrTensorType
+    : public pir::Type::
+          TypeBase<SparseCsrTensorType, pir::Type, SparseCsrTensorTypeStorage> {
+ public:
+  using Base::Base;
+
+  pir::Type dtype() const;
+  const common::DDim &dims() const;
+  common::DataLayout data_layout() const;
+  pir::DenseTensorType non_zero_crows() const;
+  pir::DenseTensorType non_zero_cols() const;
+  pir::DenseTensorType non_zero_elements() 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 SparseCsrTensorType dyn_cast_impl(pir::Type type);
+
+  static SparseCsrTensorType get(pir::IrContext *ctx,
+                                 pir::Type dtype,
+                                 const common::DDim &dims,
+                                 common::DataLayout layout,
+                                 pir::DenseTensorType non_zero_crows,
+                                 pir::DenseTensorType non_zero_cols,
+                                 pir::DenseTensorType non_zero_elements) {
+    return Base::get(ctx,
+                     dtype,
+                     dims,
+                     layout,
+                     non_zero_crows,
+                     non_zero_cols,
+                     non_zero_elements);
+  }
+};
+
 }  // 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)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SparseCsrTensorType)
diff --git a/paddle/fluid/pir/dialect/operator/ir/type_storage.h b/paddle/fluid/pir/dialect/operator/ir/type_storage.h
@@ -286,5 +286,120 @@ struct SparseCooTensorTypeStorage : public pir::TypeStorage {
   pir::DenseTensorType non_zero_elements_;
   bool coalesced_ = false;
 };
+
+struct SparseCsrTensorTypeStorage : public pir::TypeStorage {
+  ///
+  /// \brief Declare ParamKey according to parameter type.
+  ///
+  using ParamKey = std::tuple<pir::Type,
+                              common::DDim,
+                              common::DataLayout,
+                              pir::DenseTensorType,
+                              pir::DenseTensorType,
+                              pir::DenseTensorType>;
+  SparseCsrTensorTypeStorage(pir::Type dtype,
+                             common::DDim dims,
+                             common::DataLayout layout,
+                             pir::DenseTensorType non_zero_crows,
+                             pir::DenseTensorType non_zero_cols,
+                             pir::DenseTensorType non_zero_elements)
+      : dtype_(dtype),
+        dims_(dims),
+        layout_(layout),
+        non_zero_crows_(non_zero_crows),
+        non_zero_cols_(non_zero_cols),
+        non_zero_elements_(non_zero_elements) {}
+
+  ///
+  /// \brief Each derived TypeStorage must define a Construct method, which
+  /// StorageManager uses to construct a derived TypeStorage.
+  ///
+  static SparseCsrTensorTypeStorage* Construct(const ParamKey& key) {
+    return new SparseCsrTensorTypeStorage(std::get<0>(key),
+                                          std::get<1>(key),
+                                          std::get<2>(key),
+                                          std::get<3>(key),
+                                          std::get<4>(key),
+                                          std::get<5>(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 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<2>(key))));
+    // hash DenseTensorType
+    auto tuple1 = std::make_tuple(std::get<3>(key).dtype(),
+                                  std::get<3>(key).dims(),
+                                  std::get<3>(key).data_layout(),
+                                  std::get<3>(key).lod(),
+                                  std::get<3>(key).offset());
+    hash_value = pir::detail::hash_combine(
+        hash_value, DenseTensorTypeStorage::HashValue(tuple1));
+    // hash DenseTensorType
+    auto tuple2 = 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(tuple2));
+    // hash DenseTensorType
+    auto tuple3 = 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(tuple3));
+    return hash_value;
+  }
+
+  ///
+  /// \brief Each derived TypeStorage needs to overload operator==.
+  ///
+  bool operator==(const ParamKey& key) const {
+    return ParamKey(dtype_,
+                    dims_,
+                    layout_,
+                    non_zero_crows_,
+                    non_zero_cols_,
+                    non_zero_elements_) == key;
+  }
+
+  ParamKey GetAsKey() const {
+    return ParamKey(dtype_,
+                    dims_,
+                    layout_,
+                    non_zero_crows_,
+                    non_zero_cols_,
+                    non_zero_elements_);
+  }
+
+  ///
+  /// \brief SparseCsrTensorTypeStorage include six parameters: dims, dtype,
+  /// layout, non_zero_crows_,non_zero_cols_,non_zero_elements_.
+  ///
+
+  pir::Type dtype_;
+  common::DDim dims_;
+  common::DataLayout layout_;
+  pir::DenseTensorType non_zero_crows_;
+  pir::DenseTensorType non_zero_cols_;
+  pir::DenseTensorType non_zero_elements_;
+};
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/test/cpp/pir/core/type_test.cc b/test/cpp/pir/core/type_test.cc
@@ -263,6 +263,20 @@ TEST(type_test, sparse_coo) {
   pir::DenseTensorType none_zero_elements = pir::DenseTensorType::get(
       ctx, fp32_dtype, dims, data_layout, lod, offset);
   bool coalesced = false;
+  paddle::dialect::SparseCooTensorTypeStorage storage1(fp32_dtype,
+                                                       dims,
+                                                       non_zero_dims,
+                                                       data_layout,
+                                                       none_zero_indices,
+                                                       none_zero_elements,
+                                                       coalesced);
+  auto storage2 = std::make_tuple(fp32_dtype,
+                                  dims,
+                                  non_zero_dims,
+                                  data_layout,
+                                  none_zero_indices,
+                                  none_zero_elements,
+                                  coalesced);
   pir::Type pir_type =
       paddle::dialect::SparseCooTensorType::get(ctx,
                                                 fp32_dtype,
@@ -272,7 +286,7 @@ TEST(type_test, sparse_coo) {
                                                 none_zero_indices,
                                                 none_zero_elements,
                                                 coalesced);
-
+  EXPECT_TRUE(storage1 == storage2);
   EXPECT_EQ(pir_type.isa<paddle::dialect::SparseCooTensorType>(), true);
   paddle::dialect::SparseCooTensorType sparse_coo_tensor_type =
       pir_type.dyn_cast<paddle::dialect::SparseCooTensorType>();
@@ -302,6 +316,51 @@ TEST(type_test, pd_op_dialect) {
   EXPECT_EQ(select_rows_dtype.offset(), offset);
 }
 
+TEST(type_test, sparse_csr) {
+  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::DataLayout data_layout = common::DataLayout::NCHW;
+  pir::LoD lod = {{0, 1, 2}};
+  size_t offset = 0;
+  pir::DenseTensorType non_zero_crows = pir::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+  pir::DenseTensorType non_zero_cols = pir::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+  pir::DenseTensorType non_zero_elements = pir::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+  paddle::dialect::SparseCsrTensorTypeStorage storage1(fp32_dtype,
+                                                       dims,
+                                                       data_layout,
+                                                       non_zero_crows,
+                                                       non_zero_cols,
+                                                       non_zero_elements);
+  auto storage2 = std::make_tuple(fp32_dtype,
+                                  dims,
+                                  data_layout,
+                                  non_zero_crows,
+                                  non_zero_cols,
+                                  non_zero_elements);
+  pir::Type pir_type =
+      paddle::dialect::SparseCsrTensorType::get(ctx,
+                                                fp32_dtype,
+                                                dims,
+                                                data_layout,
+                                                non_zero_crows,
+                                                non_zero_cols,
+                                                non_zero_elements);
+  EXPECT_TRUE(storage1 == storage2);
+  EXPECT_EQ(pir_type.isa<paddle::dialect::SparseCsrTensorType>(), true);
+  paddle::dialect::SparseCsrTensorType sparse_csr_tensor_type =
+      pir_type.dyn_cast<paddle::dialect::SparseCsrTensorType>();
+  EXPECT_EQ(sparse_csr_tensor_type.dims(), dims);
+  EXPECT_EQ(sparse_csr_tensor_type.data_layout(), data_layout);
+  EXPECT_EQ(sparse_csr_tensor_type.non_zero_crows(), non_zero_crows);
+  EXPECT_EQ(sparse_csr_tensor_type.non_zero_cols(), non_zero_cols);
+  EXPECT_EQ(sparse_csr_tensor_type.non_zero_elements(), non_zero_elements);
+}
+
 TEST(type_test, type_util) {
   pir::IrContext *ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
