add spmd rules for greater_than and less_than op. (#68133)

zxcd · web-flow · commit 757d6af93eaa · 2024-09-24T20:59:29.000+08:00
* add spmd rules for greater_than and less_than op.

* fix

* fix

* add length check
diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc
@@ -448,6 +448,14 @@ PD_REGISTER_SPMD_RULE(
     not_equal,
     PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmd),
     PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmdReverse));
+PD_REGISTER_SPMD_RULE(
+    greater_than,
+    PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmd),
+    PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmdReverse));
+PD_REGISTER_SPMD_RULE(
+    less_than,
+    PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmd),
+    PD_INFER_SPMD(phi::distributed::ElementwiseBinaryInferSpmdReverse));
 PD_REGISTER_SPMD_RULE(swiglu,
                       PD_INFER_SPMD(phi::distributed::SwiGLUInferSpmd),
                       PD_INFER_SPMD(phi::distributed::SwiGLUInferSpmdReverse));
diff --git a/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml b/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml
@@ -187,6 +187,7 @@
   output : Tensor(out)
   infer_meta :
     func : CompareInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : greater_than
   inplace: (x -> out)
@@ -217,6 +218,7 @@
   output : Tensor(out)
   infer_meta :
     func : CompareInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : less_than
   inplace: (x -> out)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -424,6 +424,7 @@
   output : Tensor(out)
   infer_meta :
     func : CompareInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : greater_than
   data_transform :
@@ -511,6 +512,7 @@
   output : Tensor(out)
   infer_meta :
     func : CompareInferMeta
+    spmd_rule : ElementwiseBinaryInferSpmd
   kernel :
     func : less_than
   data_transform :
diff --git a/test/auto_parallel/spmd_rules/test_elementwise_rule.py b/test/auto_parallel/spmd_rules/test_elementwise_rule.py
@@ -559,6 +559,127 @@ def test_backward_multi_mesh_dim_broadcast(self):
             infered_output_dist_attrs[0].dims_mapping, [-1, 0, -1, 1]
         )
 
+    def test_single_mesh_dim_greater_than(self):
+        binary_rule = core.get_phi_spmd_rule("greater_than")
+        # [0, -1], [-1, -1] --> [0, -1], [0, -1], [0, -1]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        self.y_dist_tensor_spec.set_dims_mapping([-1, -1])
+        result_dist_attrs = binary_rule.infer_forward(
+            self.x_dist_tensor_spec, self.y_dist_tensor_spec
+        )
+        self.assertEqual(len(result_dist_attrs), 2)
+
+        infered_input_dist_attrs = result_dist_attrs[0]
+        infered_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(infered_input_dist_attrs), 2)
+        self.assertEqual(len(infered_output_dist_attrs), 1)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [0, -1])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [0, -1])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [0, -1])
+
+        # [0, -1], [-1, 0] --> [0, -1], [0, -1], [0, -1]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        self.y_dist_tensor_spec.set_dims_mapping([-1, 0])
+        result_dist_attrs = binary_rule.infer_forward(
+            self.x_dist_tensor_spec, self.y_dist_tensor_spec
+        )
+        self.assertEqual(len(result_dist_attrs), 2)
+
+        infered_input_dist_attrs = result_dist_attrs[0]
+        infered_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(infered_input_dist_attrs), 2)
+        self.assertEqual(len(infered_output_dist_attrs), 1)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [0, -1])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [0, -1])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [0, -1])
+
+    def test_single_mesh_dim_broadcast_greater_than(self):
+        binary_rule = core.get_phi_spmd_rule("greater_than")
+        self.x_dist_tensor_spec.shape = [1, 64]
+        self.y_dist_tensor_spec.shape = [64]
+
+        # [-1, 0], [-1] --> [-1, 0], [0], [-1, 0]
+        self.x_dist_tensor_spec.set_dims_mapping([-1, 0])
+        self.y_dist_tensor_spec.set_dims_mapping([-1])
+
+        resulted_dist_attrs = binary_rule.infer_forward(
+            self.x_dist_tensor_spec, self.y_dist_tensor_spec
+        )
+
+        self.assertEqual(len(resulted_dist_attrs), 2)
+
+        infered_input_dist_attrs = resulted_dist_attrs[0]
+        infered_output_dist_attrs = resulted_dist_attrs[1]
+
+        self.assertEqual(len(infered_input_dist_attrs), 2)
+        self.assertEqual(len(infered_output_dist_attrs), 1)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [-1, 0])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [0])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [-1, 0])
+
+    def test_single_mesh_dim_less_than(self):
+        binary_rule = core.get_phi_spmd_rule("less_than")
+        # [0, -1], [-1, -1] --> [0, -1], [0, -1], [0, -1]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        self.y_dist_tensor_spec.set_dims_mapping([-1, -1])
+        result_dist_attrs = binary_rule.infer_forward(
+            self.x_dist_tensor_spec, self.y_dist_tensor_spec
+        )
+        self.assertEqual(len(result_dist_attrs), 2)
+
+        infered_input_dist_attrs = result_dist_attrs[0]
+        infered_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(infered_input_dist_attrs), 2)
+        self.assertEqual(len(infered_output_dist_attrs), 1)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [0, -1])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [0, -1])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [0, -1])
+
+        # [0, -1], [-1, 0] --> [0, -1], [0, -1], [0, -1]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        self.y_dist_tensor_spec.set_dims_mapping([-1, 0])
+        result_dist_attrs = binary_rule.infer_forward(
+            self.x_dist_tensor_spec, self.y_dist_tensor_spec
+        )
+        infered_input_dist_attrs = result_dist_attrs[0]
+        infered_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [0, -1])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [0, -1])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [0, -1])
+
+    def test_single_mesh_dim_broadcast_less_than(self):
+        binary_rule = core.get_phi_spmd_rule("less_than")
+        self.x_dist_tensor_spec.shape = [1, 64]
+        self.y_dist_tensor_spec.shape = [64]
+
+        # [-1, 0], [-1] --> [-1, 0], [0], [-1, 0]
+        self.x_dist_tensor_spec.set_dims_mapping([-1, 0])
+        self.y_dist_tensor_spec.set_dims_mapping([-1])
+
+        resulted_dist_attrs = binary_rule.infer_forward(
+            self.x_dist_tensor_spec, self.y_dist_tensor_spec
+        )
+
+        self.assertEqual(len(resulted_dist_attrs), 2)
+
+        infered_input_dist_attrs = resulted_dist_attrs[0]
+        infered_output_dist_attrs = resulted_dist_attrs[1]
+
+        self.assertEqual(len(infered_input_dist_attrs), 2)
+        self.assertEqual(len(infered_output_dist_attrs), 1)
+
+        self.assertEqual(infered_input_dist_attrs[0].dims_mapping, [-1, 0])
+        self.assertEqual(infered_input_dist_attrs[1].dims_mapping, [0])
+        self.assertEqual(infered_output_dist_attrs[0].dims_mapping, [-1, 0])
+
 
 if __name__ == "__main__":
     unittest.main()
