[DistDialect] Dist Interface (#62895)

JZ-LIANG · web-flow · commit 90e62ce9d797 · 2024-03-21T11:04:59.000+08:00
* dist interface

* interface
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_interface.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_interface.cc
@@ -0,0 +1,19 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_interface.h"
+
+namespace paddle::dialect {}  // namespace paddle::dialect
+
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::DistTypeInterface)
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_interface.h b/paddle/fluid/pir/dialect/distributed/ir/dist_interface.h
@@ -0,0 +1,53 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "paddle/pir/include/core/cast_utils.h"
+#include "paddle/pir/include/core/dll_decl.h"
+#include "paddle/pir/include/core/type.h"
+
+namespace paddle {
+namespace dialect {
+
+class IR_API DistTypeInterface
+    : public pir::TypeInterfaceBase<DistTypeInterface> {
+ public:
+  struct Concept {
+    /// Defined these methods with the interface.
+    explicit Concept(pir::Type (*local_type)(pir::Type))
+        : local_type(local_type) {}
+    pir::Type (*local_type)(pir::Type);
+  };
+
+  template <class ConcreteType>
+  struct Model : public Concept {
+    static Type local_type(Type type) {
+      return pir::cast<ConcreteType>(type).local_type();
+    }
+    Model() : Concept(local_type) {}
+  };
+
+  DistTypeInterface(pir::Type type, Concept *impl)
+      : pir::TypeInterfaceBase<DistTypeInterface>(type), impl_(impl) {}
+
+  pir::Type local_type() { return impl_->local_type(*this); }
+
+ private:
+  Concept *impl_;
+};
+
+}  // namespace dialect
+}  // namespace paddle
+
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::DistTypeInterface)
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_type.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_type.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_type.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/type_storage.h"
+#include "paddle/pir/include/core/ir_context.h"
 
 namespace paddle {
 namespace dialect {
@@ -57,6 +58,15 @@ common::DDim InferLocalDDim(const common::DDim& global_ddim,
   return local_ddim;
 }
 
+auto DistDenseTensorType::local_type() const -> Type {
+  return pir::DenseTensorType::get(pir::IrContext::Instance(),
+                                   dtype(),
+                                   local_ddim(),
+                                   data_layout(),
+                                   lod(),
+                                   offset());
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_type.h b/paddle/fluid/pir/dialect/distributed/ir/dist_type.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_interface.h"
 #include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/core/type.h"
 
@@ -29,18 +30,23 @@ class DistDenseTensorType
     : public pir::Type::TypeBase<DistDenseTensorType,
                                  pir::Type,
                                  DistDenseTensorTypeStorage,
-                                 pir::WrapTypeInterface> {
+                                 pir::WrapTypeInterface,
+                                 DistTypeInterface> {
  public:
   using Base::Base;
+  using LoD = pir::DenseTensorTypeStorage::LoD;
 
   pir::DenseTensorType dense_tensor_type() const;
   TensorDistAttribute tensor_dist_attr() const;
   const common::DDim& global_ddim() const { return dense_tensor_type().dims(); }
   const common::DDim& local_ddim() const;
   Type dtype() const { return dense_tensor_type().dtype(); }
   DataLayout data_layout() const { return dense_tensor_type().data_layout(); }
+  const LoD& lod() const { return dense_tensor_type().lod(); }
+  size_t offset() const { return dense_tensor_type().offset(); }
 
   Type prim_type() { return dense_tensor_type(); }
+  Type local_type() const;
 
   ProcessMeshAttribute process_mesh_attr() const {
     return tensor_dist_attr().process_mesh_attr();
diff --git a/test/cpp/pir/distributed/dist_dialect_test.cc b/test/cpp/pir/distributed/dist_dialect_test.cc
@@ -16,6 +16,7 @@
 
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_dialect.h"
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_interface.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_op.h"
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_type.h"
 #include "paddle/fluid/pir/dialect/distributed/transforms/mix_to_dist_pass.h"
@@ -167,6 +168,53 @@ TEST(dist_dense_tensor_type_test, warp_type_interface) {
             dense_tensor_type);
 }
 
+TEST(dist_dense_tensor_type_test, dist_interface) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<DistDialect>();
+  ctx->GetOrRegisterDialect<OperatorDialect>();
+  std::vector<int64_t> mesh_shape = {2, 3};
+  std::vector<int64_t> process_ids = {0, 1, 2, 3, 4, 5};
+  std::vector<std::string> dim_names = {"x", "y"};
+  phi::distributed::ProcessMesh process_mesh(
+      mesh_shape, process_ids, dim_names);
+  auto mesh_attr = ProcessMeshAttribute::get(ctx, process_mesh);
+
+  std::vector<int64_t> dims_mapping = {0, -1};
+  paddle::flat_hash_map<int64_t, phi::ReduceType> partial_status{
+      {1, phi::ReduceType::kRedSum}};
+  // construct a TensorDistAttribute.
+  auto tensor_dist_attr =
+      TensorDistAttribute::get(ctx, mesh_attr, dims_mapping, partial_status);
+
+  pir::Type fp32_dtype = pir::Float32Type::get(ctx);
+  common::DDim dims = {4, 8};
+  common::DDim local_dims = {2, 8};
+  common::DataLayout data_layout = common::DataLayout::NCHW;
+  pir::LoD lod = {{0, 1, 2}};
+  size_t offset = 0;
+  pir::DenseTensorType dense_tensor_type = pir::DenseTensorType::get(
+      ctx, fp32_dtype, dims, data_layout, lod, offset);
+
+  pir::Type dist_densor_type =
+      DistDenseTensorType::get(ctx, dense_tensor_type, tensor_dist_attr);
+
+  EXPECT_TRUE(dist_densor_type.isa<pir::DenseTensorType>());
+  EXPECT_EQ(dist_densor_type.dyn_cast<pir::DenseTensorType>(),
+            dense_tensor_type);
+
+  // test local cast
+  auto local_dense_tensor_type = dist_densor_type.dyn_cast<DistTypeInterface>()
+                                     .local_type()
+                                     .dyn_cast<pir::DenseTensorType>();
+  EXPECT_TRUE(local_dense_tensor_type.isa<pir::DenseTensorType>());
+  EXPECT_FALSE(local_dense_tensor_type.isa<DistDenseTensorType>());
+  EXPECT_EQ(local_dense_tensor_type.dtype().isa<pir::Float32Type>(), true);
+  EXPECT_EQ(local_dense_tensor_type.dims(), local_dims);
+  EXPECT_EQ(local_dense_tensor_type.data_layout(), data_layout);
+  EXPECT_EQ(local_dense_tensor_type.lod(), lod);
+  EXPECT_EQ(local_dense_tensor_type.offset(), offset);
+}
+
 TEST(operation_dist_attr_test, base) {
   pir::IrContext* ctx = pir::IrContext::Instance();
   ctx->GetOrRegisterDialect<DistDialect>();
