1 /** 2 * Copyright 2019 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_FRONTEND_PARALLEL_OPS_INFO_MATMUL_INFO_H_ 18 #define MINDSPORE_CCSRC_FRONTEND_PARALLEL_OPS_INFO_MATMUL_INFO_H_ 19 20 #include <memory> 21 #include <string> 22 #include <unordered_map> 23 #include <vector> 24 25 #include "utils/ms_utils.h" 26 #include "ir/value.h" 27 #include "frontend/parallel/auto_parallel/operator_costmodel.h" 28 #include "frontend/parallel/ops_info/operator_info.h" 29 #include "frontend/parallel/strategy.h" 30 31 namespace mindspore { 32 namespace parallel { 33 class MatMulBase : public OperatorInfo { 34 public: MatMulBase(const std::string & name,const Shapes & inputs_shape,const Shapes & outputs_shape,const PrimitiveAttrs & attrs)35 MatMulBase(const std::string &name, const Shapes &inputs_shape, const Shapes &outputs_shape, 36 const PrimitiveAttrs &attrs) 37 : OperatorInfo(name, inputs_shape, outputs_shape, attrs, std::make_shared<MatMulCost>()) {} 38 ~MatMulBase() override = default; 39 40 Status Init(const StrategyPtr &strategy) override; 41 Status InitForCostModel(const StrategyPtr &strategy) override; 42 43 // Generate all strategies and the corresponding cost for this MatMul operator 44 Status GenerateStrategies(int64_t stage_id) override; 45 std::vector<StrategyPtr> GenerateOpStrategies(int64_t stage_id) override; 46 Status SetCostUnderStrategy(const StrategyPtr &strategy) override; 47 Status PrepareStrategy(int64_t stage_id, size_t dev_num, Dimensions combined_partitions, size_t input0_shape_size, 48 size_t input1_shape_size, StrategyPtr *sp); 49 50 Status SwapLastTwoElements(Shape *shape); 51 52 protected: 53 Status InferForwardCommunication() override; 54 Status InferTensorInfo() override; // the forward_reduce_scatter mode need to override this function 55 Status InferDevMatrixShape() override; 56 Status InferTensorMap() override; 57 Status InferTensorLayout(TensorLayouts *inputs_layout, TensorLayouts *outputs_layout); 58 void InitTensorInfoForCost(std::vector<TensorInfo> *); 59 Status CheckForTensorSliceValid() const; 60 Status GetAttrs() override; 61 62 bool transpose_a_ = false; 63 bool transpose_b_ = false; 64 bool forward_reduce_scatter_ = false; 65 int64_t field_size_ = 0; 66 size_t mat_a_dimension_ = 0; 67 size_t mat_b_dimension_ = 0; 68 Shape origin_dev_matrix_shape_; 69 }; 70 71 class MatMul : public MatMulBase { 72 public: MatMul(const std::string & name,const Shapes & inputs_shape,const Shapes & outputs_shape,const PrimitiveAttrs & attrs)73 MatMul(const std::string &name, const Shapes &inputs_shape, const Shapes &outputs_shape, const PrimitiveAttrs &attrs) 74 : MatMulBase(name, inputs_shape, outputs_shape, attrs) {} 75 ~MatMul() override = default; 76 77 protected: 78 Status CheckStrategy(const StrategyPtr &strategy) override; 79 }; 80 81 class MatMulInfo : public MatMul { 82 public: MatMulInfo(const std::string & name,const Shapes & inputs_shape,const Shapes & outputs_shape,const PrimitiveAttrs & attrs)83 MatMulInfo(const std::string &name, const Shapes &inputs_shape, const Shapes &outputs_shape, 84 const PrimitiveAttrs &attrs) 85 : MatMul(name, inputs_shape, outputs_shape, attrs) {} 86 ~MatMulInfo() override = default; 87 }; 88 89 class BatchMatMulInfo : public MatMul { 90 public: BatchMatMulInfo(const std::string & name,const Shapes & inputs_shape,const Shapes & outputs_shape,const PrimitiveAttrs & attrs)91 BatchMatMulInfo(const std::string &name, const Shapes &inputs_shape, const Shapes &outputs_shape, 92 const PrimitiveAttrs &attrs) 93 : MatMul(name, inputs_shape, outputs_shape, attrs) {} 94 ~BatchMatMulInfo() override = default; 95 96 std::shared_ptr<Strategys> GenerateBatchStrategies() override; 97 }; 98 } // namespace parallel 99 } // namespace mindspore 100 #endif // MINDSPORE_CCSRC_FRONTEND_PARALLEL_OPS_INFO_MATMUL_INFO_H_ 101