mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/strategy.py

@@ -0,0 +1,67 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# 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
+#
+# http://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.
+# ============================================================================
+"""The base class of tiling strategy"""
+from abc import ABCMeta
+from abc import abstractmethod
+from collections import namedtuple
+
+TilingPara = namedtuple("TilingPara", "Br last_Br Bc last_Bc Tr Tc")
+
+
+class TilingStrategy(metaclass=ABCMeta):
+    """Tiling strategy interface. All implementations should be defined in this module,
+    otherwise, the UT will fail.
+    """
+
+    _strategies = {}
+
+    def __init__(self, Nq, N, head_dim) -> None:
+        super().__init__()
+        self.Nq = Nq
+        self.N = N
+        self.Br = None
+        self.last_Br = None
+        self.Bc = None
+        self.last_Bc = None
+        self.Tr = None
+        self.Tc = None
+        self.d = head_dim
+
+    def __init_subclass__(cls, **kwargs):
+        TilingStrategy._strategies[cls.strategy_name()] = cls
+
+    @classmethod
+    @abstractmethod
+    def strategy_name(cls):
+        """strategy name"""
+        raise NotImplementedError
+
+    @classmethod
+    def from_strategy_name(cls, stgy_name: str):
+        """from strategy name"""
+        stgy_clz = TilingStrategy._strategies.get(stgy_name)
+        if stgy_clz is None:
+            raise Exception(f"Strategy:{stgy_name} not supported")
+
+        return stgy_clz
+
+    @abstractmethod
+    def tiling(self) -> TilingPara:
+        """tiling"""
+        raise NotImplementedError
+
+    def gen_tiling_para(self) -> TilingPara:
+        """gen tiling para"""
+        return TilingPara(self.Br, self.last_Br, self.Bc, self.last_Bc, self.Tr, self.Tc)

mindspore/ops/_op_impl/_custom_op/flash_attention/tiling_strategy/wukong_tiling.py

@@ -0,0 +1,62 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# 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
+#
+# http://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.
+# ============================================================================
+"""wukong tiling"""
+from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.strategy import TilingPara
+from mindspore.ops._op_impl._custom_op.flash_attention.tiling_strategy.strategy import TilingStrategy
+
+
+class WukongTiling(TilingStrategy):
+    """A tiling strategy implementation for wukonghuahua model shape"""
+
+    @classmethod
+    def strategy_name(cls):
+        return "wukong"
+
+    def tiling(self) -> TilingPara:
+        """
+        反向的空间分布待详细分析
+        N = (4096, 1024, 256, 64) 或 77
+        Nq = (4096, 1024, 256, 64)
+        d = dv = (40, 80, 160,  160)
+        """
+        if self.N <= 77:  # [77, 64]
+            # cross-attention or self-attention of (64, 64, 160)
+            self.Bc = self.N
+            self.Tc = self.N // self.Bc
+            if self.d <= 80:  # [40, 80]
+                # 内存瓶颈为在ub中对P*V结果[Br, dv]进行cast
+                self.Br = min(self.Nq, 64)
+                self.Tr = self.Nq // self.Br
+            else:
+                self.Br = min(self.Nq, 64)
+                self.Tr = self.Nq // self.Br
+        else:
+            # self-attention
+            if self.N == 256:
+                self.Bc = 64
+                self.Tc = 1
+                # 内存瓶颈为在ub中对Q*K的结果[Br, Bc]进行cast
+                self.Br = 64
+                self.Tr = self.Nq // self.Br
+            else:
+                self.Bc = 64
+                self.Tc = self.N // self.Bc
+                self.Br = 64
+                self.Tr = self.Nq // self.Br
+
+        self.last_Br = self.Br
+        self.last_Bc = self.Bc
+
+        return self.gen_tiling_para()

mindspore/ops/_op_impl/_custom_op/matmul_cube_dense_left_impl.py

@@ -187,8 +187,8 @@ def core(shape_a_temp, shape_b_temp, shape_output, kernel_name):
                 tik_instance.mmad(resmatmul_local_ub_local_l0c, input_1_local_l1_local_l0a,
                                   input_2_local_l1_local_l0b, 128, 128, 256, 0)
                 tik_instance.data_move(resmatmul_local_ub, resmatmul_local_ub_local_l0c, 0, 1, 128, 0, 0, 1)
-                tik_instance.data_move(resmatmul[cc6 * 256 * 1008 + core_m_idx * 8 * 256 + core_n_idx * 512 * 1008]
-                                       , resmatmul_local_ub, 0, 16, 256 // 2, 0, 55 * 16 * 2 // 2)
+                tik_instance.data_move(resmatmul[cc6 * 256 * 1008 + core_m_idx * 8 * 256 + core_n_idx * 512 * 1008],
+                                       resmatmul_local_ub, 0, 16, 256 // 2, 0, 55 * 16 * 2 // 2)
         with tik_instance.else_scope():
             tik_instance.data_move(input_1_local_l1, input_x1[core_m_idx * (8 * 256 + 128 * 1008)], 0, 7, 112,
                                    56 * 16, 0)

mindspore/ops/_op_impl/aicpu/__init__.py

@@ -14,12 +14,16 @@
 
 """aicpu ops"""
 from .adaptive_max_pool_3d_grad import _adaptive_max_pool_3d_grad_aicpu
+from .adaptive_max_pool_2d import _adaptive_max_pool_2d_aicpu
 from .adaptive_max_pool_2d_grad import _adaptive_max_pool_2d_grad_aicpu
 from .adaptive_avg_pool_3d_grad import _adaptiveavgpool3d_grad_aicpu
 from .adaptive_avg_pool_3d import _adaptiveavgpool3d_aicpu
+from .adaptive_max_pool_3d import _adaptive_max_pool_3d_aicpu
 from .tile import _tile_aicpu
 from .tanh import _tanh_aicpu
 from .less import _less_aicpu
+from .lstsq import _lstsq_aicpu
+from .left_shift import _left_shift_aicpu
 from .add import _add_aicpu
 from .sparse_matrix_transpose import _sparse_matrix_transpose_aicpu
 from .sparse_matrix_nnz import _sparse_matrix_nnz_aicpu
@@ -46,6 +50,7 @@ from .arg_min import _arg_min_aicpu
 from .argmin_with_value import _argmin_with_value_aicpu
 from .avgpool_v1 import _avgpool_v1_aicpu
 from .avgpool_grad_v1 import _avgpool_grad_v1_aicpu
+from .matrix_logarithm import _matrix_logarithm_aicpu
 from .matrix_solve import _matrix_solve_aicpu
 from .betainc import _betainc_aicpu
 from .bartlett_window import _bartlett_window_aicpu
@@ -91,6 +96,8 @@ from .acos_grad import _acos_grad_aicpu
 from .expand import _expand_aicpu
 from .expand_dims import _expand_dims_aicpu
 from .randperm import _randperm_aicpu
+from .randperm_v2 import _randperm_v2_aicpu
+from .random_poisson import _random_poisson_aicpu
 from .random_choice_with_mask import _random_choice_with_mask_aicpu
 from .rsqrt import _rsqrt_aicpu
 from .sqrt import _sqrt_aicpu
@@ -101,11 +108,14 @@ from .search_sorted import _search_sorted_aicpu
 from .stack import _stack_aicpu
 from .unstack import _unstack_aicpu
 from .unsorted_segment_sum import _unsorted_segment_sum_aicpu
+from .unsorted_segment_prod import _unsorted_segment_prod_aicpu
 from .addcmul import _addcmul_aicpu
 from .uniform_candidate_sampler import _uniform_candidate_sampler_aicpu
 from .log_uniform_candidate_sampler import _log_uniform_candidate_sampler_aicpu
 from .compute_accidental_hits import _compute_accidental_hits_aicpu
 from .ctcloss import _ctcloss_aicpu
+from .ctc_loss_v2 import _ctc_loss_v2_aicpu
+from .ctc_loss_v2_grad import _ctc_loss_v2_grad_aicpu
 from .reverse_sequence import _reverse_sequence_aicpu
 from .log_matrix_determinant import _log_matrix_determinant_aicpu
 from .crop_and_resize import _crop_and_resize_aicpu
@@ -119,6 +129,7 @@ from .tanh_grad import _tanh_grad_aicpu
 from .cast import _cast_aicpu
 from .mirror_pad import _mirror_pad_aicpu
 from .mirror_pad_grad import _mirror_pad_grad_aicpu
+from .masked_scatter import _masked_scatter_aicpu
 from .masked_select import _masked_select_aicpu
 from .masked_select_grad import _masked_select_grad_aicpu
 from .mul import _mul_aicpu
@@ -129,8 +140,13 @@ from .sub import _sub_aicpu
 from .not_equal import _not_equal_aicpu
 from .poisson import _poisson_aicpu
 from .update_cache import _update_cache_aicpu
+from .upsample_nearest_3d import _upsample_nearest_3d_aicpu
+from .upsample_nearest_3d_grad import _upsample_nearest_3d_grad_aicpu
+from .upsample_trilinear_3d import _upsample_trilinear_3d_aicpu
+from .upsample_trilinear_3d_grad import _upsample_trilinear_3d_grad_aicpu
 from .upper_bound import _upper_bound_aicpu
 from .cache_swap_table import _cache_swap_table_aicpu
+from .uniform import _uniform_aicpu
 from .uniform_int import _uniform_int_aicpu
 from .uniform_real import _uniform_real_aicpu
 from .standard_laplace import _standard_laplace_aicpu
@@ -142,12 +158,13 @@ from .fused_sparse_adam import _fused_sparse_adam_aicpu
 from .fused_sparse_lazy_adam import _fused_sparse_lazy_adam_aicpu
 from .fused_sparse_ftrl import _fused_sparse_ftrl_aicpu
 from .sparse_fill_empty_rows_grad import _sparse_fill_empty_rows_grad_aicpu
+from .sparse_reorder import _sparse_reorder_aicpu
 from .sparse_reshape import _sparse_reshape_aicpu
 from .sparse_segment_sqrt_n_grad import _sparse_segment_sqrt_n_grad_aicpu
 from .sparse_segment_sum import _sparse_segment_sum_aicpu
 from .sparse_segment_sum_with_num_segments import _sparse_segment_sum_with_num_segments_aicpu
 from .sparse_softmax_cross_entropy_with_logits_v2 import _sparse_softmax_cross_entropy_with_logits_v2_aicpu
-from .sparsesparsemaximum import _sparsesparsemaximum_aicpu
+from .sparse_sparse_maximum import _sparse_sparse_maximum_aicpu
 from .split import _split_aicpu
 from .transpose import _transpose_aicpu
 from .tril_indices import _tril_indices_aicpu
@@ -164,6 +181,8 @@ from .stack_push_pop import _stack_push_aicpu
 from .stack_push_pop import _stack_pop_aicpu
 from .asinh import _asinh_aicpu
 from .stack_push_pop import _stack_destroy_aicpu
+from .matrix_band_part import _matrix_band_part_aicpu
+from .matrix_exp import _matrix_exp_aicpu
 from .matrix_diag_v3 import _matrix_diag_v3_aicpu
 from .matrix_diag_part_v3 import _matrix_diag_part_v3_aicpu
 from .tan import _tan_aicpu
@@ -189,6 +208,7 @@ from .environ_get import _environ_get_aicpu
 from .environ_destroy_all import _environ_destroy_all_aicpu
 from .cross import _cross_aicpu
 from .check_numerics import _check_numerics_aicpu
+from .cummax import _cummax_aicpu
 from .cumsum import _cumsum_aicpu
 from .round import _round_aicpu
 from .stft import _stft_aicpu
@@ -213,6 +233,7 @@ from .scatter_nd_update import _scatter_nd_update_aicpu
 from .scatter_nd_max import _scatter_nd_max_aicpu
 from .conj import _conj_aicpu
 from .scatter_nd_min import _scatter_nd_min_aicpu
+from .scatter_add_with_axis import _scatter_add_with_axis_aicpu
 from .compare_and_bitpack import _compare_and_bitpack_aicpu
 from .addcdiv import _addcdiv_aicpu
 from .unique_consecutive import _unique_consecutive_aicpu
@@ -226,6 +247,7 @@ from .reservoir_replay_buffer import _rrb_sample_op_cpu
 from .reservoir_replay_buffer import _rrb_destroy_op_cpu
 from .concat_offset import _concat_offset_aicpu
 from .range import _range_aicpu
+from .range_v2 import _range_v2_aicpu
 from .slice_grad import _slice_grad_aicpu
 from .median import _median_aicpu
 from .median_grad import _median_grad_aicpu
@@ -233,6 +255,7 @@ from .reduce_sum import _reduce_sum_aicpu
 from .adaptive_avg_pool_2d import _adaptive_avg_pool_2d_aicpu
 from .adaptive_avg_pool_2d_grad import _adaptive_avg_pool_2d_grad_aicpu
 from .fill_v2 import _fill_v2_aicpu
+from .fill_diagonal import _fill_diagonal_aicpu
 from .data_format_vec_permute import _data_format_vec_permute_aicpu
 from .multinomial import _multinomial_aicpu
 from .fft_with_size import _fft_with_size_aicpu
@@ -254,6 +277,7 @@ from .complex import _complex_aicpu
 from .complex_abs import _complex_abs_aicpu
 from .concat import _concat_aicpu
 from .cos import _cos_aicpu
+from .count_nonzero import _count_nonzero_aicpu
 from .csr_sparse_matrix_to_dense import _csr_sparse_matrix_to_dense_aicpu
 from .cumprod import _cumprod_aicpu
 from .exp import _exp_aicpu
@@ -269,6 +293,7 @@ from .one_hot import _one_hot_aicpu
 from .orgqr import _orgqr_aicpu
 from .parameterized_truncated_normal import _parameterized_truncated_normal_aicpu
 from .polar import _polar_aicpu
+from .polygamma import _polygamma_aicpu
 from .pdist_grad import _pdist_grad_aicpu
 from .ragged_range import _raggedrange_aicpu
 from .ragged_tensor_to_sparse import _ragged_tensor_to_sparse_aicpu
@@ -294,12 +319,14 @@ from .cumulative_logsumexp import _cumulative_logsumexp_aicpu
 from .sparse_segment_sqrt_n import _sparse_segment_sqrt_n_aicpu
 from .scale_and_translate import _scale_and_translate_aicpu
 from .quant_dtype_cast import _quant_dtype_cast_aicpu
+from .quantile import _quantile_aicpu
 from .fse_decode import _fse_decode_aicpu
 from .dense_to_csr_sparse_matrix import _dense_to_csr_sparse_matrix_aicpu
 from .dense_to_sparse_set_operation import _dense_to_sparse_set_operation_aicpu
 from .diag import _diag_aicpu
 from .diagonal import _diagonal_aicpu
 from .diag_part import _diag_part_aicpu
+from .digamma import _digamma_aicpu
 from .bias_add import _bias_add_aicpu
 from .bias_add_grad import _bias_add_grad_aicpu
 from .eig import _eig_aicpu
@@ -318,6 +345,8 @@ from .heaviside import _heaviside_aicpu
 from .hypot import _hypot_aicpu
 from .identity_n import _identity_n_aicpu
 from .index_fill import _index_fill_aicpu
+from .index_put import _index_put_aicpu
+from .inplace_index_add import _inplace_index_add_aicpu
 from .kldivloss import _kldiv_loss_aicpu
 from .kldivlossgrad import _kldiv_loss_grad_aicpu
 from .lcm import _lcm_aicpu
@@ -333,7 +362,9 @@ from .pad_v3 import _pad_v3_aicpu
 from .cholesky import _cholesky_aicpu
 from .hsv_to_rgb import _hsv_to_rgb_aicpu
 from .im2col import _im2col_aicpu
+from .bessel_i0 import _bessel_i0_aicpu
 from .lu_solve import _lu_solve_aicpu
+from .lu import _lu_aicpu
 from .relu_grad_v3 import _relu_grad_v3_aicpu
 from .resize_bicubic import _resize_bicubic_aicpu
 from .extract_glimpse import _extract_glimpse_aicpu
@@ -359,6 +390,7 @@ from .layer_norm_grad_grad import _layernorm_grad_grad_aicpu
 from .list_diff import _list_diff_aicpu
 from .log import _log_aicpu
 from .logspace import _logspace_aicpu
+from .lgamma import _lgamma_aicpu
 from .matrix_inverse import _matrix_inverse_aicpu
 from .matrix_power import _matrix_power_aicpu
 from .max_pool3d_grad_with_argmax import _max_pool3d_grad_with_argmax_aicpu
@@ -375,6 +407,9 @@ from .non_deterministic_ints import _non_deterministic_ints_aicpu
 from .pow import _pow_aicpu
 from .real import _real_aicpu
 from .resize_area import _resize_area_aicpu
+from .segment_mean import _segment_mean_aicpu
+from .segment_min import _segment_min_aicpu
+from .segment_prod import _segment_prod_aicpu
 from .segment_sum import _segment_sum_aicpu
 from .set_size import _set_size_aicpu
 from .slice import _slice_aicpu
@@ -386,6 +421,7 @@ from .sparse_tensor_dense_mat_mul import _sparse_tensor_dense_mat_mul_aicpu
 from .trace import _trace_aicpu
 from .tracegrad import _tracegrad_aicpu
 from .tridiagonal_solve import _tridiagonal_solve_aicpu
+from .tridiagonal_matmul import _tridiagonal_matmul_aicpu
 from .truncated_normal import _truncated_normal_aicpu
 from .glu import _glu_aicpu
 from .deformable_offsets import _deformable_offsets_aicpu
@@ -397,4 +433,8 @@ from .bernoulli import _bernoulli_aicpu
 from .glu_grad import _glu_grad_aicpu
 from .sspaddmm import _sspaddmm_aicpu
 from .sequence_addn import _sequence_addn_aicpu
+from .sequence_concat import _sequence_concat_aicpu
+from .sequence_stack import _sequence_stack_aicpu
 from .affine_grid import _affine_grid_aicpu
+from .depth_to_space import _depth_to_space_aicpu
+from .eps import _eps_aicpu

mindspore/ops/_op_impl/aicpu/adaptive_max_pool_2d.py

@@ -0,0 +1,37 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# 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
+#
+# http://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.
+# ============================================================================
+
+"""AdaptiveMaxPool2D op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+adaptive_max_pool_2d_op_info = AiCPURegOp("AdaptiveMaxPool2D")                                            \
+    .fusion_type("OPAQUE")                                                                                \
+    .attr("output_size", "listInt")                                                                       \
+    .input(0, "x", "required")                                                                            \
+    .output(0, "y", "required")                                                                           \
+    .output(1, "argmax", "required")                                                                      \
+    .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.I64_Default) \
+    .get_op_info()
+
+
+@op_info_register(adaptive_max_pool_2d_op_info)
+def _adaptive_max_pool_2d_aicpu():
+    """AdaptiveMaxPool2D aicpu register"""
+    return

mindspore/ops/_op_impl/aicpu/bias_add_grad.py

@@ -31,7 +31,6 @@ bias_add_grad_op_info = AiCPURegOp("BiasAddGrad") \
     .dtype_format(DataType.I64_Default, DataType.I64_Default) \
     .dtype_format(DataType.F16_Default, DataType.F16_Default) \
     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
-    .dtype_format(DataType.F64_Default, DataType.F64_Default) \
     .dtype_format(DataType.C64_Default, DataType.C64_Default) \
     .dtype_format(DataType.C128_Default, DataType.C128_Default) \
     .get_op_info()

mindspore/ops/_op_impl/aicpu/cast.py

@@ -32,6 +32,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.U8_Default, DataType.F32_Default) \
     .dtype_format(DataType.U8_Default, DataType.F64_Default) \
     .dtype_format(DataType.U8_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.U8_Default, DataType.C64_Default) \
+    .dtype_format(DataType.U8_Default, DataType.C128_Default) \
     .dtype_format(DataType.U16_Default, DataType.U8_Default) \
     .dtype_format(DataType.U16_Default, DataType.U16_Default) \
     .dtype_format(DataType.U16_Default, DataType.U32_Default) \
@@ -44,6 +46,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.U16_Default, DataType.F32_Default) \
     .dtype_format(DataType.U16_Default, DataType.F64_Default) \
     .dtype_format(DataType.U16_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.U16_Default, DataType.C64_Default) \
+    .dtype_format(DataType.U16_Default, DataType.C128_Default) \
     .dtype_format(DataType.U32_Default, DataType.U8_Default) \
     .dtype_format(DataType.U32_Default, DataType.U16_Default) \
     .dtype_format(DataType.U32_Default, DataType.U32_Default) \
@@ -56,6 +60,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.U32_Default, DataType.F32_Default) \
     .dtype_format(DataType.U32_Default, DataType.F64_Default) \
     .dtype_format(DataType.U32_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.U32_Default, DataType.C64_Default) \
+    .dtype_format(DataType.U32_Default, DataType.C128_Default) \
     .dtype_format(DataType.U64_Default, DataType.U8_Default) \
     .dtype_format(DataType.U64_Default, DataType.U16_Default) \
     .dtype_format(DataType.U64_Default, DataType.U32_Default) \
@@ -68,6 +74,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.U64_Default, DataType.F32_Default) \
     .dtype_format(DataType.U64_Default, DataType.F64_Default) \
     .dtype_format(DataType.U64_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.U64_Default, DataType.C64_Default) \
+    .dtype_format(DataType.U64_Default, DataType.C128_Default) \
     .dtype_format(DataType.I8_Default, DataType.U8_Default) \
     .dtype_format(DataType.I8_Default, DataType.U16_Default) \
     .dtype_format(DataType.I8_Default, DataType.U32_Default) \
@@ -80,6 +88,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.I8_Default, DataType.F32_Default) \
     .dtype_format(DataType.I8_Default, DataType.F64_Default) \
     .dtype_format(DataType.I8_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.I8_Default, DataType.C64_Default) \
+    .dtype_format(DataType.I8_Default, DataType.C128_Default) \
     .dtype_format(DataType.I16_Default, DataType.U8_Default) \
     .dtype_format(DataType.I16_Default, DataType.U16_Default) \
     .dtype_format(DataType.I16_Default, DataType.U32_Default) \
@@ -92,6 +102,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.I16_Default, DataType.F32_Default) \
     .dtype_format(DataType.I16_Default, DataType.F64_Default) \
     .dtype_format(DataType.I16_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.I16_Default, DataType.C64_Default) \
+    .dtype_format(DataType.I16_Default, DataType.C128_Default) \
     .dtype_format(DataType.I32_Default, DataType.U8_Default) \
     .dtype_format(DataType.I32_Default, DataType.U16_Default) \
     .dtype_format(DataType.I32_Default, DataType.U32_Default) \
@@ -104,6 +116,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.I32_Default, DataType.F32_Default) \
     .dtype_format(DataType.I32_Default, DataType.F64_Default) \
     .dtype_format(DataType.I32_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.I32_Default, DataType.C64_Default) \
+    .dtype_format(DataType.I32_Default, DataType.C128_Default) \
     .dtype_format(DataType.I32_5HD, DataType.I64_5HD) \
     .dtype_format(DataType.I64_Default, DataType.U8_Default) \
     .dtype_format(DataType.I64_Default, DataType.U16_Default) \
@@ -117,6 +131,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.I64_Default, DataType.F32_Default) \
     .dtype_format(DataType.I64_Default, DataType.F64_Default) \
     .dtype_format(DataType.I64_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.I64_Default, DataType.C64_Default) \
+    .dtype_format(DataType.I64_Default, DataType.C128_Default) \
     .dtype_format(DataType.F16_Default, DataType.U8_Default) \
     .dtype_format(DataType.F16_Default, DataType.U16_Default) \
     .dtype_format(DataType.F16_Default, DataType.U32_Default) \
@@ -129,6 +145,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.F16_Default, DataType.F32_Default) \
     .dtype_format(DataType.F16_Default, DataType.F64_Default) \
     .dtype_format(DataType.F16_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.F16_Default, DataType.C64_Default) \
+    .dtype_format(DataType.F16_Default, DataType.C128_Default) \
     .dtype_format(DataType.F32_Default, DataType.U8_Default) \
     .dtype_format(DataType.F32_Default, DataType.U16_Default) \
     .dtype_format(DataType.F32_Default, DataType.U32_Default) \
@@ -141,6 +159,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.F32_Default, DataType.F32_Default) \
     .dtype_format(DataType.F32_Default, DataType.F64_Default) \
     .dtype_format(DataType.F32_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.F32_Default, DataType.C64_Default) \
+    .dtype_format(DataType.F32_Default, DataType.C128_Default) \
     .dtype_format(DataType.F64_Default, DataType.U8_Default) \
     .dtype_format(DataType.F64_Default, DataType.U16_Default) \
     .dtype_format(DataType.F64_Default, DataType.U32_Default) \
@@ -153,6 +173,8 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.F64_Default, DataType.F32_Default) \
     .dtype_format(DataType.F64_Default, DataType.F64_Default) \
     .dtype_format(DataType.F64_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.F64_Default, DataType.C64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.C128_Default) \
     .dtype_format(DataType.BOOL_Default, DataType.U8_Default) \
     .dtype_format(DataType.BOOL_Default, DataType.U16_Default) \
     .dtype_format(DataType.BOOL_Default, DataType.U32_Default) \
@@ -165,6 +187,36 @@ cast_op_info = AiCPURegOp("Cast") \
     .dtype_format(DataType.BOOL_Default, DataType.F32_Default) \
     .dtype_format(DataType.BOOL_Default, DataType.F64_Default) \
     .dtype_format(DataType.BOOL_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.BOOL_Default, DataType.C64_Default) \
+    .dtype_format(DataType.BOOL_Default, DataType.C128_Default) \
+    .dtype_format(DataType.C64_Default, DataType.U8_Default) \
+    .dtype_format(DataType.C64_Default, DataType.U16_Default) \
+    .dtype_format(DataType.C64_Default, DataType.U32_Default) \
+    .dtype_format(DataType.C64_Default, DataType.U64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I8_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I16_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I32_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.F16_Default) \
+    .dtype_format(DataType.C64_Default, DataType.F32_Default) \
+    .dtype_format(DataType.C64_Default, DataType.F64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.C64_Default, DataType.C64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.C128_Default) \
+    .dtype_format(DataType.C128_Default, DataType.U8_Default) \
+    .dtype_format(DataType.C128_Default, DataType.U16_Default) \
+    .dtype_format(DataType.C128_Default, DataType.U32_Default) \
+    .dtype_format(DataType.C128_Default, DataType.U64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I8_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I16_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I32_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.F16_Default) \
+    .dtype_format(DataType.C128_Default, DataType.F32_Default) \
+    .dtype_format(DataType.C128_Default, DataType.F64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.BOOL_Default) \
+    .dtype_format(DataType.C128_Default, DataType.C64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.C128_Default) \
     .get_op_info()
 
 @op_info_register(cast_op_info)

mindspore/ops/_op_impl/aicpu/coalesce.py

@@ -23,6 +23,8 @@ coalesce_op_info = AiCPURegOp("Coalesce") \
     .output(0, "y_indices", "required") \
     .output(1, "y_values", "required") \
     .output(2, "y_shape", "required") \
+    .dtype_format(DataType.I64_Default, DataType.F64_Default, DataType.I64_Default, DataType.I64_Default,
+                  DataType.F64_Default, DataType.I64_Default) \
     .dtype_format(DataType.I64_Default, DataType.F32_Default, DataType.I64_Default, DataType.I64_Default,
                   DataType.F32_Default, DataType.I64_Default) \
     .dtype_format(DataType.I64_Default, DataType.F16_Default, DataType.I64_Default, DataType.I64_Default,

mindspore/ops/_op_impl/aicpu/col2im.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-
 """Col2Im op"""
 from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
 
@@ -27,6 +26,9 @@ col2im_op_info = AiCPURegOp("Col2Im") \
     .attr("padding", "listInt") \
     .dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default) \
     .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.F64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I32_Default, DataType.C64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I32_Default, DataType.C128_Default) \
     .get_op_info()
 
 

mindspore/ops/_op_impl/aicpu/count_nonzero.py

@@ -0,0 +1,43 @@
+# Copyright 2022 Huawei Technologies Co., Ltd
+#
+# 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
+#
+# http://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.
+# ============================================================================
+
+"""CountNonZero op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+count_nonzero_op_info = AiCPURegOp("CountNonZero") \
+    .fusion_type("OPAQUE") \
+    .input(0, "x", "required") \
+    .output(0, "y", "required") \
+    .attr("dims", "listInt")\
+    .dtype_format(DataType.I8_Default, DataType.I64_Default) \
+    .dtype_format(DataType.I16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.I32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.I64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U8_Default, DataType.I64_Default)  \
+    .dtype_format(DataType.U16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.U64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F16_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F32_Default, DataType.I64_Default) \
+    .dtype_format(DataType.F64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.C64_Default, DataType.I64_Default) \
+    .dtype_format(DataType.C128_Default, DataType.I64_Default) \
+    .get_op_info()
+
+
+@op_info_register(count_nonzero_op_info)
+def _count_nonzero_aicpu():
+    """CountNonZero AiCPU register"""
+    return

mindspore/ops/_op_impl/aicpu/dropout_genmask.py

@@ -25,8 +25,14 @@ dropout_genmask_op_info = AiCPURegOp("DropoutGenMask") \
     .output(0, "y", "required") \
     .attr("Seed0", "int") \
     .attr("Seed1", "int") \
+    .dtype_format(DataType.I64_Default, DataType.F16_Default, DataType.U64_Default, DataType.U64_Default,
+                  DataType.U8_Default) \
     .dtype_format(DataType.I32_Default, DataType.F16_Default, DataType.U64_Default, DataType.U64_Default,
                   DataType.U8_Default) \
+    .dtype_format(DataType.I64_Default, DataType.F32_Default, DataType.U64_Default, DataType.U64_Default,
+                  DataType.U8_Default) \
+    .dtype_format(DataType.I32_Default, DataType.F32_Default, DataType.U64_Default, DataType.U64_Default,
+                  DataType.U8_Default) \
     .get_op_info()
 
 @op_info_register(dropout_genmask_op_info)

mindspore/ops/_op_impl/aicpu/eps.py

@@ -0,0 +1,32 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2023. 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
+#
+# http://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.
+# ============================================================================
+
+"""Eps op"""
+from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
+
+eps_op_info = AiCPURegOp("Eps") \
+    .fusion_type("OPAQUE") \
+    .input(0, "x", "required") \
+    .output(0, "y", "required") \
+    .dtype_format(DataType.F16_Default, DataType.F16_Default)   \
+    .dtype_format(DataType.F32_Default, DataType.F32_Default)   \
+    .dtype_format(DataType.F64_Default, DataType.F64_Default)   \
+    .get_op_info()
+
+
+@op_info_register(eps_op_info)
+def _eps_aicpu():
+    """Eps AiCPU register"""
+    return