dist2dense_pass fix shape errors in shard randomly sampled data (#68067)

* dist2dense_pass fix shape errors in shard randomly sampled data

* add unit test case

* fix coverage issues

Author: jeff41404

paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc

@@ -103,6 +103,35 @@ void ProcessDistBlock(pir::Block* block) {
         new_dims.push_back(pir::Int64Attribute::get(ctx, local_dims[index]));
       }
       prev_op->set_attribute("value", pir::ArrayAttribute::get(ctx, new_dims));
+    } else if (op_item->isa<RandintOp>() || op_item->isa<GaussianOp>() ||
+               op_item->isa<UniformOp>()) {
+      auto local_dims =
+          op_item->result_type(0).dyn_cast<pir::DenseTensorType>().dims();
+      auto shape_value = op_item->operand_source(0);
+      auto prev_op = shape_value.defining_op();
+      PADDLE_ENFORCE((prev_op != nullptr),
+                     common::errors::PreconditionNotMet(
+                         "The shape of randint, gaussian and uniform mush be "
+                         "the result of "
+                         "FullIntArrayOp, not null"));
+      PADDLE_ENFORCE_EQ(
+          prev_op->isa<FullIntArrayOp>(),
+          true,
+          common::errors::PreconditionNotMet(
+              "The shape of randint, gaussian and uniform mush be the result "
+              "of FullIntArrayOp."));
+      auto array_attr = prev_op->attribute<pir::ArrayAttribute>("value");
+      PADDLE_ENFORCE_EQ(
+          array_attr.size(),
+          local_dims.size(),
+          common::errors::PreconditionNotMet(
+              "The shape of randint, gaussian and uniform element's size must "
+              "equal to result's dim size."));
+      std::vector<pir::Attribute> new_dims;
+      for (int index = 0; index < local_dims.size(); ++index) {
+        new_dims.push_back(pir::Int64Attribute::get(ctx, local_dims[index]));
+      }
+      prev_op->set_attribute("value", pir::ArrayAttribute::get(ctx, new_dims));
     }
     // TODO(2024-Q2) not all op are dist type
     // PADDLE_ENFORCE_EQ(

test/auto_parallel/pir/test_static_pir_program.py

@@ -16,6 +16,7 @@
 
 import paddle
 import paddle.distributed as dist
+from paddle.base.libpaddle.pir import apply_dist2dense_pass
 from paddle.distributed.auto_parallel.static.mix_to_dist_pass import (
     apply_mix2dist_pass,
 )
@@ -228,6 +229,67 @@ def test_build_with_apply_mix2dist_pass(self):
         self.assertEqual(input2_shape.dist_attr().process_mesh, mesh)
         self.assertEqual(input2_shape.dist_attr().dims_mapping, [-1])
 
+    def test_build_with_apply_dist2dense_pass(self):
+        paddle.enable_static()
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.base.Program()
+            with paddle.base.program_guard(main_program):
+                mesh = dist.ProcessMesh([0, 1], dim_names=['dp'])
+                input1 = paddle.randint(low=0, high=1000, shape=[8, 4])
+                output1 = dist.shard_tensor(input1, mesh, [dist.Shard(0)])
+
+                input2 = paddle.randn([4, 8])
+                output2 = dist.shard_tensor(input2, mesh, [dist.Shard(1)])
+
+                self.assertTrue(input1.is_dense_tensor_type())
+                self.assertTrue(input2.is_dense_tensor_type())
+
+        self.assertTrue(main_program.num_ops() == 6)
+
+        self.assertFalse(input1.use_empty())
+        self.assertFalse(input2.use_empty())
+
+        self.assertTrue(output1.use_empty())
+        self.assertTrue(output2.use_empty())
+
+        self.assertFalse(input1.get_defining_op().has_attr("op_dist_attr"))
+        self.assertFalse(input2.get_defining_op().has_attr("op_dist_attr"))
+
+        # check dist type
+        self.assertTrue(output1.is_dist_dense_tensor_type())
+        self.assertTrue(output2.is_dist_dense_tensor_type())
+
+        # run apply_mix2dist_pass and apply_dist2dense_pass
+        apply_mix2dist_pass(main_program)
+        apply_dist2dense_pass(main_program)
+
+        # after apply_mix2dist_pass, the program changed
+        # and after apply_dist2dense_pass, the operator in program do not have dist_attr
+        self.assertTrue(main_program.num_ops() == 4)
+
+        self.assertTrue(input1.is_dense_tensor_type())
+        self.assertTrue(input2.is_dense_tensor_type())
+
+        self.assertFalse(input1.get_defining_op().has_attr("op_dist_attr"))
+        self.assertFalse(input2.get_defining_op().has_attr("op_dist_attr"))
+
+        # check shape attribute of full_int_array op
+        input1_shape = input1.get_defining_op().operand_source(0)
+        input1_shape_op = input1_shape.get_defining_op()
+        self.assertFalse(input1_shape_op.has_attr("op_dist_attr"))
+        input1_shape_op_attr = input1_shape_op.attrs()
+        self.assertEqual(input1_shape_op_attr['value'], [4, 4])
+
+        input2_shape = input2.get_defining_op().operand_source(0)
+        input2_shape_op = input2_shape.get_defining_op()
+        self.assertFalse(input2_shape_op.has_attr("op_dist_attr"))
+        input2_shape_op_attr = input2_shape_op.attrs()
+        self.assertEqual(input2_shape_op_attr['value'], [4, 4])
+
+        # check shape of input1 and input2
+        self.assertEqual(input1.shape, [4, 4])
+        self.assertEqual(input2.shape, [4, 4])
+
 
 if __name__ == "__main__":
     unittest.main()