!245 Add bool type check in communication operator

Merge pull request !245 from chentingting/add_bool_type_check_in_comm_op

mindspore/ccsrc/parallel/auto_parallel/edge_costmodel.cc

@@ -61,11 +61,12 @@ Status Edge::InitEdgeCost() {
       auto target_output_lyt = target_output.second[prev_op_output_index_].tensor_layout();
       auto target_output_str = target_output.first;
       auto type_length = prev_op_->GetOutputTypeLengths()[prev_op_output_index_];
+      auto type = prev_op_->outputs_type()[prev_op_output_index_];
       for (auto& target_input : next_op_input_) {
         auto target_input_lyt = target_input.second[next_op_input_index_].tensor_layout();
         auto target_input_str = target_input.first;
         CostPtr cost;
-        if (GetRedistributionCost(target_output_lyt, target_input_lyt, type_length, &cost) != SUCCESS) {
+        if (GetRedistributionCost(target_output_lyt, target_input_lyt, type_length, type, &cost) != SUCCESS) {
           MS_LOG(EXCEPTION) << "Failure: redistribution cost calculation failed";
         }
         MS_EXCEPTION_IF_NULL(cost);
@@ -99,7 +100,7 @@ Status Edge::InitEdgeCost() {
 }
 
 Status Edge::GetRedistributionCost(const TensorLayout& prev_op_output_layout, const TensorLayout& next_op_input_layout,
-                                   size_t type_length, CostPtr* cost) {
+                                   size_t type_length, TypePtr type, CostPtr* cost) {
   MS_EXCEPTION_IF_NULL(prev_op_);
   MS_EXCEPTION_IF_NULL(cost);
   RankList dev_list = prev_op_->global_device_list();
@@ -119,6 +120,13 @@ Status Edge::GetRedistributionCost(const TensorLayout& prev_op_output_layout, co
   double backward_comm_cost = tensor_redistribution.backward_comm_cost();
   double computation_cost = tensor_redistribution.computation_cost();
 
+  // Now AllGather, ReduceScatter, AlltoAll don't support bool type
+  MS_EXCEPTION_IF_NULL(type);
+  if ((type->type_id() == kNumberTypeBool) && (comm_cost > 0)) {
+    computation_cost = INF;
+    comm_cost = INF;
+    MS_LOG(WARNING) << "Communication Operators don't support bool dtype!";
+  }
   *cost = std::make_shared<Cost>(type_length * computation_cost, type_length * comm_cost);
   (*cost)->communication_without_parameter_ = type_length * comm_cost;
   (*cost)->communication_with_partial_para_ =

mindspore/ccsrc/parallel/auto_parallel/edge_costmodel.h

@@ -84,7 +84,7 @@ class Edge {
   // and the input tensor layout of v, return the redistribution cost,
   // and the op_list to carry out the redistribution.
   Status GetRedistributionCost(const TensorLayout& prev_op_output_layout, const TensorLayout& next_op_input_layout,
-                               size_t, CostPtr* cost);
+                               size_t, TypePtr type, CostPtr* cost);
 
   void set_pre_op_output(const std::vector<std::pair<std::shared_ptr<Strategy>, std::vector<TensorInfo>>>& output_set) {
     pre_op_output_ = output_set;

mindspore/ccsrc/parallel/ops_info/operator_info.cc

@@ -1197,6 +1197,16 @@ Status OperatorInfo::SetInputAndOutputTypeLength(const std::vector<size_t>& inpu
   return SUCCESS;
 }
 
+Status OperatorInfo::set_outputs_type(const std::vector<TypePtr>& outputs_type) {
+  if (outputs_type.size() != outputs_shape_.size()) {
+    MS_LOG(ERROR) << "Outputs type: " << outputs_type.size()
+                  << " do not have the same number of outputs shape: " << outputs_shape_.size();
+    return FAILED;
+  }
+  outputs_type_ = outputs_type;
+  return SUCCESS;
+}
+
 void OperatorInfo::BreakingTiesForPerferringDataParallel(const StrategyPtr& stra, const CostPtr& cost) {
   if (!stra->GetInputDim().empty() && !stra->GetInputDim()[0].empty()) {
     CheckGlobalDeviceManager();

mindspore/ccsrc/parallel/ops_info/operator_info.h

@@ -60,7 +60,8 @@ class OperatorInfo {
         outputs_shape_(std::move(outputs_shape)),
         attrs_(std::move(attrs)),
         is_alive_(true),
-        cost_(cost) {
+        cost_(cost),
+        outputs_type_() {
     std::vector<bool> not_parameteter(inputs_shape_.size(), false);
     is_parameter_ = not_parameteter;
     refkey_parameter_name_ = "";
@@ -71,6 +72,11 @@ class OperatorInfo {
   Status set_is_parameter(const std::vector<bool>& is_parameter);
   Status SetInputAndOutputTypeLength(const std::vector<size_t>& input_lengths,
                                      const std::vector<size_t>& output_lengths);
+  // Set outputs dtype.
+  // If only one output, outputs_type.size() is 1.
+  // If output is tuple, outputs_type.size() is greater than 1.
+  Status set_outputs_type(const std::vector<TypePtr>& outputs_type);
+  const std::vector<TypePtr>& outputs_type() const { return outputs_type_; }
   virtual Status Init(const StrategyPtr& strategy) = 0;
   virtual Status InitForCostModel(const StrategyPtr& strategy) = 0;  // only init the necessary parts
 
@@ -229,6 +235,7 @@ class OperatorInfo {
 
  private:
   OperatorCostPtr cost_;
+  std::vector<TypePtr> outputs_type_;
 };
 
 Shape GetSliceShape(const Shape& tensor_shape, const Dimensions& strategy);

mindspore/ccsrc/parallel/ops_info/ops_utils.h

@@ -42,6 +42,7 @@ constexpr size_t SoftmaxCrossEntropyWithLogitsAttrSize = 1;
 constexpr size_t SoftmaxCrossEntropyWithLogitsInputsSize = 2;
 constexpr size_t SoftmaxCrossEntropyWithLogitsOutputsSize = 2;
 constexpr double EPS = 1e-6;
+constexpr double INF = 1e20;
 
 constexpr char AUTO_PARALLEL_RUN_ONCE_ONLY[] = "auto_parallel_run_once_only";
 constexpr char SEMI_AUTO_PARALLEL_RUN_ONCE_ONLY[] = "semi_auto_parallel_run_once_only";

mindspore/ccsrc/parallel/step_auto_parallel.cc

@@ -255,12 +255,9 @@ size_t GetInputsTypeLen(const AnfNodePtr &input) {
   return input_type_len;
 }
 
-// Given the node, return the element length of input and output
+std::vector<size_t> ExtractInputTypeLengthByNode(const CNodePtr &node) {
-std::vector<std::vector<size_t>> ExtractInputAndOutputTypeLengthByNode(const CNodePtr &node) {
   MS_EXCEPTION_IF_NULL(node);
   std::vector<size_t> inputs_type_len;
-  std::vector<size_t> outputs_type_len;
-  std::vector<std::vector<size_t>> all_types;
   std::vector<AnfNodePtr> node_inputs{node->inputs()};
 
   // extract input element length
@@ -278,9 +275,13 @@ std::vector<std::vector<size_t>> ExtractInputAndOutputTypeLengthByNode(const CNo
       inputs_type_len.push_back(GetInputsTypeLen(input));
     }
   }
-  all_types.push_back(inputs_type_len);
+  return inputs_type_len;
+}
 
-  // extract output element length
+std::vector<TypePtr> ExtractOutputTypeByNode(const CNodePtr &node) {
+  MS_EXCEPTION_IF_NULL(node);
+  std::vector<TypePtr> outputs_type;
+  // extract output element type
   auto primary_output_type = node->Type();
   MS_EXCEPTION_IF_NULL(primary_output_type);
   if (primary_output_type->isa<mindspore::Tuple>()) {
@@ -290,7 +291,7 @@ std::vector<std::vector<size_t>> ExtractInputAndOutputTypeLengthByNode(const CNo
     for (auto &ele : elements) {
       if (ele->isa<mindspore::TensorType>()) {
         auto ele_element_type = ele->cast<mindspore::TensorTypePtr>()->element();
-        outputs_type_len.push_back(GetLengthOfDataType(ele_element_type));
+        outputs_type.push_back(ele_element_type);
       } else {
         MS_LOG(EXCEPTION) << "Unknown type: " << primary_output_type->type_name();
       }
@@ -299,14 +300,12 @@ std::vector<std::vector<size_t>> ExtractInputAndOutputTypeLengthByNode(const CNo
     // in this case, the output is a single tensor
     if (primary_output_type->isa<mindspore::TensorType>()) {
       auto element_type = primary_output_type->cast<mindspore::TensorTypePtr>()->element();
-      outputs_type_len.push_back(GetLengthOfDataType(element_type));
+      outputs_type.push_back(element_type);
     } else {
       MS_LOG(EXCEPTION) << "Unknown type: " << primary_output_type->type_name();
     }
   }
-  all_types.push_back(outputs_type_len);
+  return outputs_type;
-
-  return all_types;
 }
 
 // Be careful the argument is cnode_full_name, not the op_name
@@ -367,11 +366,20 @@ OperatorInfoPtr CreateTheOperatorInfo(const PrimitivePtr &prim, const CNodePtr &
     return nullptr;
   }
   // Set the data type for inputs and outputs of this OperatorInfo
-  std::vector<std::vector<size_t>> type_lengths = ExtractInputAndOutputTypeLengthByNode(cnode);
+  auto inputs_type_length = ExtractInputTypeLengthByNode(cnode);
-  if (operator_info->SetInputAndOutputTypeLength(type_lengths[0], type_lengths[1]) != SUCCESS) {
+  auto outputs_type = ExtractOutputTypeByNode(cnode);
+  std::vector<size_t> outputs_type_length;
+  outputs_type_length.reserve(outputs_type.size());
+  std::transform(outputs_type.begin(), outputs_type.end(), std::back_inserter(outputs_type_length),
+                 GetLengthOfDataType);
+  if (operator_info->SetInputAndOutputTypeLength(inputs_type_length, outputs_type_length) != SUCCESS) {
     MS_LOG(ERROR) << "Setting the lengths of inputs and outputs failed for operator: " << operator_info->name();
     return nullptr;
   }
+  if (operator_info->set_outputs_type(outputs_type) != SUCCESS) {
+    MS_LOG(ERROR) << "Setting the types of outputs failed for operator: " << operator_info->name();
+    return nullptr;
+  }
   // When the 'inputs' contains numerical values for some operators, these values should be extracted from
   // ANF graph
   auto &inputs = cnode->inputs();

mindspore/ccsrc/parallel/step_auto_parallel.h

@@ -39,7 +39,9 @@ size_t GetLengthOfDataType(const TypePtr &type);
 
 std::vector<bool> ExtractInputParameterByNode(const CNodePtr &node);
 
-std::vector<std::vector<size_t>> ExtractInputAndOutputTypeLengthByNode(const CNodePtr &node);
+std::vector<size_t> ExtractInputTypeLengthByNode(const CNodePtr &node);
+
+std::vector<TypePtr> ExtractOutputTypeByNode(const CNodePtr &node);
 
 Status ConstructCostGraphNodes(const std::vector<AnfNodePtr> &all_nodes, const FuncGraphPtr &root);
 

mindspore/ops/operations/comm_ops.py

@@ -162,6 +162,8 @@ class AllGather(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_dtype):
+        if x_dtype == mstype.bool_:
+            raise TypeError("AllGather does not support 'Bool' as the dtype of input!")
         return x_dtype
 
     def __call__(self, tensor):
@@ -219,6 +221,8 @@ class ReduceScatter(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_dtype):
+        if x_dtype == mstype.bool_:
+            raise TypeError("ReduceScatter does not support 'Bool' as the dtype of input!")
         return x_dtype
 
     def __call__(self, tensor):
@@ -276,6 +280,8 @@ class Broadcast(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_dtype):
+        if x_dtype == mstype.bool_:
+            raise TypeError("Broadcast does not support 'Bool' as the dtype of input!")
         return x_dtype
 
 
@@ -318,6 +324,8 @@ class _AlltoAll(PrimitiveWithInfer):
         return x_shape
 
     def infer_dtype(self, x_dtype):
+        if x_dtype == mstype.bool_:
+            raise TypeError("AlltoAll does not support 'Bool' as the dtype of input!")
         return x_dtype
 
     def __call__(self, tensor):

tests/ut/cpp/parallel/auto_parallel/edge_costmodel_test.cc

@@ -65,6 +65,7 @@ void TestEdgeCostModel::SetUp() {
   Shapes inputs_shape_1 = {{8, 16}, {16, 32}};
   Shapes outputs_shape_1 = {{8, 32}};
   matmul1 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_1, outputs_shape_1, attr_1);
+  matmul1->set_outputs_type({kFloat32});
 
   // matmul2
   ValuePtr transpose_a_2 = MakeValue(false);
@@ -73,6 +74,7 @@ void TestEdgeCostModel::SetUp() {
   Shapes inputs_shape_2 = {{8, 32}, {32, 16}};
   Shapes outputs_shape_2 = {{8, 16}};
   matmul2 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_2, outputs_shape_2, attr_2);
+  matmul2->set_outputs_type({kFloat32});
 
   // matmul3
   ValuePtr transpose_a_3 = MakeValue(false);
@@ -81,6 +83,7 @@ void TestEdgeCostModel::SetUp() {
   Shapes inputs_shape_3 = {{16, 8}, {8, 32}};
   Shapes outputs_shape_3 = {{16, 32}};
   matmul3 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_3, outputs_shape_3, attr_3);
+  matmul3->set_outputs_type({kFloat32});
 
   // matmul4
   ValuePtr transpose_a_4 = MakeValue(false);
@@ -89,6 +92,7 @@ void TestEdgeCostModel::SetUp() {
   Shapes inputs_shape_4 = {{8, 16}, {16, 32}};
   Shapes outputs_shape_4 = {{8, 32}};
   matmul4 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_4, outputs_shape_4, attr_4);
+  matmul4->set_outputs_type({kFloat32});
 
   // matmul5
   ValuePtr transpose_a_5 = MakeValue(false);
@@ -97,6 +101,7 @@ void TestEdgeCostModel::SetUp() {
   Shapes inputs_shape_5 = {{8, 32}, {8, 32}};
   Shapes outputs_shape_5 = {{8, 8}};
   matmul5 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_5, outputs_shape_5, attr_5);
+  matmul5->set_outputs_type({kFloat32});
 }
 
 TEST_F(TestEdgeCostModel, test_InitEdgeCost) {

tests/ut/cpp/parallel/auto_parallel/graph_costmodel_test.cc

@@ -76,6 +76,7 @@ void TestCostGraph::SetUp() {
   Shapes inputs_shape_0 = {{32, 16}, {16, 16}};
   Shapes outputs_shape_0 = {{32, 16}};
   matmul0 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_0, outputs_shape_0, attr_0);
+  matmul0->set_outputs_type({kFloat32});
 
   // matmul1
   ValuePtr transpose_a_1 = MakeValue(false);
@@ -84,6 +85,7 @@ void TestCostGraph::SetUp() {
   Shapes inputs_shape_1 = {{8, 16}, {16, 32}};
   Shapes outputs_shape_1 = {{8, 32}};
   matmul1 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_1, outputs_shape_1, attr_1);
+  matmul1->set_outputs_type({kFloat32});
 
   // matmul2
   ValuePtr transpose_a_2 = MakeValue(false);
@@ -92,6 +94,7 @@ void TestCostGraph::SetUp() {
   Shapes inputs_shape_2 = {{8, 32}, {32, 16}};
   Shapes outputs_shape_2 = {{8, 16}};
   matmul2 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_2, outputs_shape_2, attr_2);
+  matmul2->set_outputs_type({kFloat32});
 
   // matmul3
   ValuePtr transpose_a_3 = MakeValue(false);
@@ -100,6 +103,7 @@ void TestCostGraph::SetUp() {
   Shapes inputs_shape_3 = {{16, 8}, {8, 32}};
   Shapes outputs_shape_3 = {{16, 32}};
   matmul3 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_3, outputs_shape_3, attr_3);
+  matmul3->set_outputs_type({kFloat32});
 
   // matmul4
   ValuePtr transpose_a_4 = MakeValue(false);
@@ -108,6 +112,7 @@ void TestCostGraph::SetUp() {
   Shapes inputs_shape_4 = {{8, 16}, {16, 32}};
   Shapes outputs_shape_4 = {{8, 32}};
   matmul4 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_4, outputs_shape_4, attr_4);
+  matmul4->set_outputs_type({kFloat32});
 
   // matmul5
   ValuePtr transpose_a_5 = MakeValue(false);
@@ -116,6 +121,7 @@ void TestCostGraph::SetUp() {
   Shapes inputs_shape_5 = {{8, 32}, {8, 32}};
   Shapes outputs_shape_5 = {{8, 8}};
   matmul5 = std::make_shared<MatMulInfo>("matmul_info", inputs_shape_5, outputs_shape_5, attr_5);
+  matmul5->set_outputs_type({kFloat32});
 }
 
 void TestCostGraph::ConstructStarGraph2() {

tests/ut/python/communication/test_comm.py

@@ -55,7 +55,7 @@ class BroadCastNet(nn.Cell):
         self.broadcast = Broadcast(0)
 
     def construct(self, x):
-        x, = self.broadcast((x,))
+        x = self.broadcast((x))
         x = self.dense(x)
         return x