# ============================================================================ """StridedSlice op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType strided_slice_d_op_info = TBERegOp("StridedSlice") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("strided_slice_d.so") \ .compute_cost(10) \ .kernel_name("strided_slice_d") \ .partial_flag(True) \ .attr("begin", "optional", "listInt", "all") \ .attr("end", "optional", "listInt", "all") \ .attr("strides", "optional", "listInt", "all") \ .attr("begin_mask", "required", "int", "all") \ .attr("end_mask", "required", "int", "all") \ .attr("ellipsis_mask", "required", "int", "all") \ .attr("new_axis_mask", "required", "int", "all") \ .attr("shrink_axis_mask", "required", "int", "all") \ .input(0, "x", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.BOOL_Default, DataType.BOOL_Default) \ .dtype_format(DataType.I8_Default, DataType.I8_Default) \ .dtype_format(DataType.U8_Default, DataType.U8_Default) \ .dtype_format(DataType.I32_Default, DataType.I32_Default) \ .dtype_format(DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F32_Default, DataType.F32_Default) \ .get_op_info() @op_info_register(strided_slice_d_op_info) def _strided_slice_d_tbe():
from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType layer_norm_x_backprop_op_info = TBERegOp("LayerNormXBackprop") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("layer_norm_x_backprop.so") \ .compute_cost(10) \ .kernel_name("layer_norm_x_backprop") \ .partial_flag(True) \ .input(0, "dy", False, "required", "all") \ .input(1, "x", False, "required", "all") \ .input(2, "variance", False, "required", "all") \ .input(3, "mean", False, "required", "all") \ .input(4, "gamma", False, "required", "all") \ .output(0, "pd_x", False, "required", "all") \ .op_pattern("dynamicFormat") \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F16_FracNZ, DataType.F16_FracNZ, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_FracNZ) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .dtype_format(DataType.F32_FracNZ, DataType.F32_FracNZ, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_FracNZ) \ .get_op_info()
# 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. # ============================================================================ """ReduceMin op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType reduce_min_op_info = TBERegOp("ReduceMin") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("reduce_min_d.so") \ .compute_cost(10) \ .kernel_name("reduce_min_d") \ .partial_flag(True) \ .attr("axis", "required", "listInt", "all") \ .attr("keep_dims", "required", "bool", "all") \ .input(0, "x", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .op_pattern("reduce") \ .dtype_format(DataType.I8_None, DataType.I8_None) \ .dtype_format(DataType.U8_None, DataType.U8_None) \ .dtype_format(DataType.F16_None, DataType.F16_None) \ .dtype_format(DataType.F32_None, DataType.F32_None) \ .get_op_info() @op_info_register(reduce_min_op_info) def _reduce_min_tbe(): """ReduceMin TBE register""" return
# # 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. # ============================================================================ """Mul op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType mul_ds_op_info = TBERegOp("Mul") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("mul.so") \ .compute_cost(10) \ .kernel_name("mul") \ .partial_flag(True) \ .dynamic_shape(True) \ .input(0, "x", False, "required", "all") \ .input(1, "y", False, "required", "all") \ .output(0, "output", False, "required", "all") \ .op_pattern("dynamicFormat") \ .dtype_format(DataType.None_None, DataType.None_None, DataType.None_None) \ .get_op_info() @op_info_register(mul_ds_op_info) def _mul_ds_tbe(): """Mul TBE register""" return
from te.platform.fusion_manager import fusion_manager from topi import generic from topi.cce import util from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType SHAPE_SIZE_LIMIT = 2147483648 correction_mul_grad_op_info = TBERegOp("CorrectionMulGrad") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("correction_mul_grad.so") \ .compute_cost(10) \ .kernel_name("correction_mul_grad") \ .partial_flag(True) \ .op_pattern("formatAgnostic") \ .attr("channel_axis", "optional", "int", "all") \ .input(0, "dout", None, "required", None) \ .input(1, "x", None, "required", None) \ .input(2, "batch_std", None, "required", None) \ .input(3, "running_std", None, "required", None) \ .output(0, "dx", True, "required", "all") \ .output(1, "d_batch_std", True, "required", "all") \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(correction_mul_grad_op_info) def _correction_mul_grad_tbe(): """CorrectionMulGrad TBE register""" return
# 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. # ============================================================================ """Softplus op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType softplus_op_info = TBERegOp("Softplus") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("softplus.so") \ .compute_cost(10) \ .kernel_name("softplus") \ .partial_flag(True) \ .op_pattern("formatAgnostic") \ .input(0, "x", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(softplus_op_info) def _softplus_tbe(): """Softplus TBE register""" return
# 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. # ============================================================================ """MaxPool3DGrad op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType max_pool3d_grad_op_info = TBERegOp("MaxPool3DGrad") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("max_pool3d_grad.so") \ .compute_cost(10) \ .kernel_name("max_pool3d_grad") \ .partial_flag(True) \ .attr("kernel_size", "required", "listInt", "all") \ .attr("strides", "required", "listInt", "all") \ .attr("pad_list", "required", "listInt", "all") \ .attr("format", "optional", "str", "all") \ .input(0, "orig_x", False, "required", "all") \ .input(1, "orig_y", False, "required", "all") \ .input(2, "grads", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_NDC1HWC0, DataType.F16_NDC1HWC0, DataType.F16_NDC1HWC0, DataType.F32_NDC1HWC0) \ .get_op_info() @op_info_register(max_pool3d_grad_op_info) def _max_pool_3d_grad_tbe(): """MaxPool3DGrad TBE register""" return
parallel_concat_op_info = TBERegOp("ParallelConcat") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("parallel_concat.so") \ .compute_cost(10) \ .kernel_name("parallel_concat") \ .partial_flag(True) \ .attr("shape", "required", "listInt", "all") \ .attr("N", "required", "int", "all") \ .input(0, "values", False, "dynamic", "all") \ .output(0, "output_data", False, "required", "all") \ .dtype_format(DataType.BOOL_Default, DataType.BOOL_Default) \ .dtype_format(DataType.BOOL_5HD, DataType.BOOL_5HD) \ .dtype_format(DataType.I8_Default, DataType.I8_Default) \ .dtype_format(DataType.I8_5HD, DataType.I8_5HD) \ .dtype_format(DataType.U8_Default, DataType.U8_Default) \ .dtype_format(DataType.U8_5HD, DataType.U8_5HD) \ .dtype_format(DataType.I16_Default, DataType.I16_Default) \ .dtype_format(DataType.I16_5HD, DataType.I16_5HD) \ .dtype_format(DataType.U16_Default, DataType.U16_Default) \ .dtype_format(DataType.U16_5HD, DataType.U16_5HD) \ .dtype_format(DataType.I32_Default, DataType.I32_Default) \ .dtype_format(DataType.I32_5HD, DataType.I32_5HD) \ .dtype_format(DataType.U32_Default, DataType.U32_Default) \ .dtype_format(DataType.U32_5HD, DataType.U32_5HD) \ .dtype_format(DataType.I64_Default, DataType.I64_Default) \ .dtype_format(DataType.I64_5HD, DataType.I64_5HD) \ .dtype_format(DataType.U64_Default, DataType.U64_Default) \ .dtype_format(DataType.U64_5HD, DataType.U64_5HD) \ .dtype_format(DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD) \ .dtype_format(DataType.BOOL_NHWC, DataType.BOOL_NHWC) \ .dtype_format(DataType.BOOL_NCHW, DataType.BOOL_NCHW) \ .dtype_format(DataType.I8_NHWC, DataType.I8_NHWC) \ .dtype_format(DataType.I8_NCHW, DataType.I8_NCHW) \ .dtype_format(DataType.U8_NHWC, DataType.U8_NHWC) \ .dtype_format(DataType.U8_NCHW, DataType.U8_NCHW) \ .dtype_format(DataType.I16_NHWC, DataType.I16_NHWC) \ .dtype_format(DataType.I16_NCHW, DataType.I16_NCHW) \ .dtype_format(DataType.U16_NHWC, DataType.U16_NHWC) \ .dtype_format(DataType.U16_NCHW, DataType.U16_NCHW) \ .dtype_format(DataType.I32_NHWC, DataType.I32_NHWC) \ .dtype_format(DataType.I32_NCHW, DataType.I32_NCHW) \ .dtype_format(DataType.U32_NHWC, DataType.U32_NHWC) \ .dtype_format(DataType.U32_NCHW, DataType.U32_NCHW) \ .dtype_format(DataType.I64_NHWC, DataType.I64_NHWC) \ .dtype_format(DataType.I64_NCHW, DataType.I64_NCHW) \ .dtype_format(DataType.U64_NHWC, DataType.U64_NHWC) \ .dtype_format(DataType.U64_NCHW, DataType.U64_NCHW) \ .dtype_format(DataType.F16_NHWC, DataType.F16_NHWC) \ .dtype_format(DataType.F16_NCHW, DataType.F16_NCHW) \ .dtype_format(DataType.F32_NHWC, DataType.F32_NHWC) \ .dtype_format(DataType.F32_NCHW, DataType.F32_NCHW) \ .get_op_info()
apply_adam_op_info = TBERegOp("Adam") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("apply_adam.so") \ .compute_cost(10) \ .kernel_name("apply_adam") \ .partial_flag(True) \ .attr("use_locking", "optional", "bool", "true,false", "false") \ .attr("use_nesterov", "optional", "bool", "true,false", "false") \ .input(0, "var", False, "required", "all") \ .input(1, "m", False, "required", "all") \ .input(2, "v", False, "required", "all") \ .input(3, "beta1_power", False, "required", "all") \ .input(4, "beta2_power", False, "required", "all") \ .input(5, "lr", False, "required", "all") \ .input(6, "beta1", False, "required", "all") \ .input(7, "beta2", False, "required", "all") \ .input(8, "epsilon", False, "required", "all") \ .input(9, "grad", False, "required", "all") \ .output(0, "var", False, "required", "all") \ .output(1, "m", False, "required", "all") \ .output(2, "v", False, "required", "all") \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F16_FracZ, DataType.F16_FracZ, DataType.F16_FracZ, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_FracZ, DataType.F16_FracZ, DataType.F16_FracZ, DataType.F16_FracZ) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .dtype_format(DataType.F32_FracZ, DataType.F32_FracZ, DataType.F32_FracZ, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_FracZ, DataType.F32_FracZ, DataType.F32_FracZ, DataType.F32_FracZ) \ .get_op_info()
# 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. # ============================================================================ """Sub op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType sub_op_info = TBERegOp("Sub") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("sub.so") \ .compute_cost(10) \ .kernel_name("sub") \ .partial_flag(True) \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .op_pattern("broadcast") \ .dtype_format(DataType.I32_None, DataType.I32_None, DataType.I32_None) \ .dtype_format(DataType.F16_None, DataType.F16_None, DataType.F16_None) \ .dtype_format(DataType.F32_None, DataType.F32_None, DataType.F32_None) \ .get_op_info() @op_info_register(sub_op_info) def _sub_tbe(): """Add TBE register""" return
# 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. # ============================================================================ """ConfusionTransposeD op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType confusion_transpose_d_op_info = TBERegOp("ConfusionTransposeD") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("confusion_transpose_d.so") \ .compute_cost(10) \ .kernel_name("confusion_transpose_d") \ .partial_flag(True) \ .attr("perm", "required", "listInt", "all") \ .attr("shape", "required", "listInt", "all") \ .attr("transpose_first", "required", "bool", "all") \ .input(0, "x", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .op_pattern("dynamicFormat") \ .dtype_format(DataType.None_None, DataType.None_None) \ .get_op_info() @op_info_register(confusion_transpose_d_op_info) def _confusion_transpose_d_tbe(): """ConfusionTransposeD TBE register""" return
# 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. # ============================================================================ """BinaryCrossEntropy op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType binary_cross_entropy_op_info = TBERegOp("BinaryCrossEntropy") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("binary_cross_entropy.so") \ .compute_cost(10) \ .kernel_name("binary_cross_entropy") \ .partial_flag(True) \ .attr("reduction", "optional", "str", "all") \ .input(0, "x", False, "required", "all") \ .input(1, "y", False, "required", "all") \ .input(2, "weight", False, "optional", "all") \ .output(0, "output", False, "required", "all") \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(binary_cross_entropy_op_info) def _binary_cross_entropy_tbe(): """BinaryCrossEntropy TBE register""" return
pack_op_info = TBERegOp("Pack") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("pack.so") \ .compute_cost(10) \ .kernel_name("pack") \ .partial_flag(True) \ .need_check_supported(True) \ .attr("axis", "optional", "int", "all") \ .input(0, "x", False, "dynamic", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.I8_Default, DataType.I8_Default) \ .dtype_format(DataType.I16_Default, DataType.I16_Default) \ .dtype_format(DataType.I32_Default, DataType.I32_Default) \ .dtype_format(DataType.I64_Default, DataType.I64_Default) \ .dtype_format(DataType.U8_Default, DataType.U8_Default) \ .dtype_format(DataType.U16_Default, DataType.U16_Default) \ .dtype_format(DataType.U32_Default, DataType.U32_Default) \ .dtype_format(DataType.U64_Default, DataType.U64_Default) \ .dtype_format(DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.BOOL_Default, DataType.BOOL_Default) \ .dtype_format(DataType.I8_NDHWC, DataType.I8_NDHWC) \ .dtype_format(DataType.I16_NDHWC, DataType.I16_NDHWC) \ .dtype_format(DataType.I32_NDHWC, DataType.I32_NDHWC) \ .dtype_format(DataType.I64_NDHWC, DataType.I64_NDHWC) \ .dtype_format(DataType.U8_NDHWC, DataType.U8_NDHWC) \ .dtype_format(DataType.U16_NDHWC, DataType.U16_NDHWC) \ .dtype_format(DataType.U32_NDHWC, DataType.U32_NDHWC) \ .dtype_format(DataType.U64_NDHWC, DataType.U64_NDHWC) \ .dtype_format(DataType.F16_NDHWC, DataType.F16_NDHWC) \ .dtype_format(DataType.F32_NDHWC, DataType.F32_NDHWC) \ .dtype_format(DataType.BOOL_NDHWC, DataType.BOOL_NDHWC) \ .get_op_info()
from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType sparse_apply_ftrl_d_op_info = TBERegOp("SparseApplyFtrl") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("sparse_apply_ftrl_d.so") \ .compute_cost(10) \ .kernel_name("sparse_apply_ftrl_d") \ .partial_flag(True) \ .attr("lr", "required", "float", "all") \ .attr("l1", "required", "float", "all") \ .attr("l2", "required", "float", "all") \ .attr("lr_power", "required", "float", "all") \ .attr("use_locking", "optional", "bool", "true,false", "false") \ .input(0, "var", False, "required", "all") \ .input(1, "accum", False, "required", "all") \ .input(2, "linear", False, "required", "all") \ .input(3, "grad", False, "required", "all") \ .input(4, "indices", False, "required", "all") \ .output(0, "var", False, "required", "all") \ .output(1, "accum", False, "required", "all") \ .output(2, "linear", False, "required", "all") \ .dtype_format(DataType.F32_NCHW, DataType.F32_NCHW, DataType.F32_NCHW, DataType.F32_NCHW, DataType.I32_NCHW, DataType.F32_NCHW, DataType.F32_NCHW, DataType.F32_NCHW) \ .dtype_format(DataType.F32_NHWC, DataType.F32_NHWC, DataType.F32_NHWC, DataType.F32_NHWC, DataType.I32_NHWC, DataType.F32_NHWC, DataType.F32_NHWC, DataType.F32_NHWC) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.I32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .get_op_info()
from te import tvm from te.platform.fusion_manager import fusion_manager from topi import generic from topi.cce import util from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType fake_quant_per_layer_op_info = TBERegOp("FakeQuantPerLayer") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("fake_quant_per_layer.so") \ .compute_cost(10) \ .kernel_name("fake_quant_per_layer") \ .partial_flag(True) \ .attr("symmetric", "optional", "bool", "all") \ .attr("narrow_range", "optional", "bool", "all") \ .attr("num_bits", "optional", "int", "all") \ .input(0, "x", None, "required", None) \ .input(1, "min", None, "required", None) \ .input(2, "max", None, "required", None) \ .output(0, "y", True, "required", "all") \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(fake_quant_per_layer_op_info) def _fake_quant_per_layer_tbe(): """FakeQuantPerLayer TBE register""" return
# 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. # ============================================================================ """MaxPoolGradWithArgmax op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType max_pool_grad_with_argmax_op_info = TBERegOp("MaxPoolGradWithArgmax") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("max_pool_grad_with_argmax.so") \ .compute_cost(10) \ .kernel_name("max_pool_grad_with_argmax") \ .partial_flag(True) \ .attr("ksize", "required", "listInt", "all") \ .attr("strides", "required", "listInt", "all") \ .attr("padding", "required", "str", "all") \ .input(0, "x", False, "required", "all") \ .input(1, "grad", False, "required", "all") \ .input(2, "argmax", False, "optional", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.U16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.I64_5HD, DataType.F16_5HD) \ .get_op_info() @op_info_register(max_pool_grad_with_argmax_op_info) def _max_pool_grad_with_argmax_tbe(): """MaxPoolGradWithArgmax TBE register""" return
# # 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. # ============================================================================ """BatchNorm op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType bn_training_reduce_op_info = TBERegOp("BNTrainingReduce") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("bn_training_reduce.so") \ .compute_cost(10) \ .kernel_name("bn_training_reduce") \ .partial_flag(True) \ .input(0, "x", False, "required", "all", reshape_type="NC") \ .output(0, "sum", False, "required", "all") \ .output(1, "square_sum", False, "required", "all") \ .is_dynamic_format(True) \ .dtype_format(DataType.F16_None, DataType.F32_None, DataType.F32_None) \ .dtype_format(DataType.F32_None, DataType.F32_None, DataType.F32_None) \ .get_op_info() @op_info_register(bn_training_reduce_op_info) def _bn_training_reduce_tbe(): """BNTrainingReduce TBE register""" return
# 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. # ============================================================================ """Sin op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType sin_op_info = TBERegOp("Sin") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("sin.so") \ .compute_cost(10) \ .kernel_name("sin") \ .partial_flag(True) \ .op_pattern("formatAgnostic") \ .input(0, "x", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_None, DataType.F16_None) \ .dtype_format(DataType.F32_None, DataType.F32_None) \ .get_op_info() @op_info_register(sin_op_info) def _sin_tbe(): """Sin TBE register""" return
# 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. # ============================================================================ """SigmoidGrad op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType sigmoid_cross_entropy_with_logits_op_info = TBERegOp("SigmoidGrad") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("sigmoid_grad.so") \ .compute_cost(10) \ .kernel_name("sigmoid_grad") \ .partial_flag(True) \ .input(0, "x", False, "required", "all") \ .input(1, "y", False, "required", "all") \ .output(0, "z", False, "required", "all") \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(sigmoid_cross_entropy_with_logits_op_info) def _sigmoid_grad_tbe(): """SigmoidGrad TBE register""" return
concat_op_info = TBERegOp("Concat") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("concat_d.so") \ .compute_cost(10) \ .kernel_name("concat_d") \ .partial_flag(True) \ .attr("axis", "required", "int", "all") \ .input(0, "input_values", False, "dynamic", "all") \ .output(0, "output_data", False, "required", "all") \ .op_pattern("dynamicFormat") \ .dtype_format(DataType.BOOL_Default, DataType.BOOL_Default) \ .dtype_format(DataType.BOOL_5HD, DataType.BOOL_5HD) \ .dtype_format(DataType.I8_Default, DataType.I8_Default) \ .dtype_format(DataType.I8_5HD, DataType.I8_5HD) \ .dtype_format(DataType.U8_Default, DataType.U8_Default) \ .dtype_format(DataType.U8_5HD, DataType.U8_5HD) \ .dtype_format(DataType.I16_Default, DataType.I16_Default) \ .dtype_format(DataType.I16_5HD, DataType.I16_5HD) \ .dtype_format(DataType.U16_Default, DataType.U16_Default) \ .dtype_format(DataType.U16_5HD, DataType.U16_5HD) \ .dtype_format(DataType.I32_Default, DataType.I32_Default) \ .dtype_format(DataType.I32_5HD, DataType.I32_5HD) \ .dtype_format(DataType.U32_Default, DataType.U32_Default) \ .dtype_format(DataType.U32_5HD, DataType.U32_5HD) \ .dtype_format(DataType.I64_Default, DataType.I64_Default) \ .dtype_format(DataType.I64_5HD, DataType.I64_5HD) \ .dtype_format(DataType.U64_Default, DataType.U64_Default) \ .dtype_format(DataType.U64_5HD, DataType.U64_5HD) \ .dtype_format(DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info()
# See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ """BatchNormGrad op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType bn_training_update_grad_op_info = TBERegOp("BNTrainingUpdateGrad") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("bn_training_update_grad.so") \ .compute_cost(10) \ .kernel_name("bn_training_update_grad") \ .partial_flag(True) \ .attr("epsilon", "optional", "float", "all") \ .input(0, "grads", False, "required", "all", reshape_type="NC") \ .input(1, "x", False, "required", "all", reshape_type="NC") \ .input(2, "batch_mean", False, "required", "all") \ .input(3, "batch_variance", False, "required", "all") \ .output(0, "diff_scale", False, "required", "all") \ .output(1, "diff_offset", False, "required", "all") \ .op_pattern("dynamicFormat") \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(bn_training_update_grad_op_info) def _bn_training_update_grad_tbe(): """BNTrainingUpdateGrad TBE register""" return
# 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. # ============================================================================ """LogicalAnd op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType logical_and_op_info = TBERegOp("LogicalAnd") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("logical_and.so") \ .compute_cost(10) \ .kernel_name("logical_and") \ .partial_flag(True) \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .output(0, "y", True, "required", "all") \ .op_pattern("broadcast") \ .dtype_format(DataType.BOOL_None, DataType.BOOL_None, DataType.BOOL_None) \ .get_op_info() @op_info_register(logical_and_op_info) def _logical_and_tbe(): """LogicalAnd TBE register""" return
# See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ """Add op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType tensor_add_op_info = TBERegOp("Add") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("add.so") \ .compute_cost(10) \ .kernel_name("add") \ .partial_flag(True) \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .is_dynamic_format(True) \ .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .dtype_format(DataType.I32_5HD, DataType.I32_5HD, DataType.I32_5HD) \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .get_op_info() @op_info_register(tensor_add_op_info) def _tensor_add_tbe(): """Add TBE register""" return
# See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ """LessEqual op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType less_equal_op_info = TBERegOp("LessEqual") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("less_equal.so") \ .compute_cost(10) \ .kernel_name("less_equal") \ .partial_flag(True) \ .attr("begin_norm_axis", "required", "int", "all") \ .attr("begin_params_axis", "required", "int", "all") \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .op_pattern("broadcast") \ .dtype_format(DataType.I8_None, DataType.I8_None, DataType.BOOL_None) \ .dtype_format(DataType.U8_None, DataType.U8_None, DataType.BOOL_None) \ .dtype_format(DataType.I32_None, DataType.I32_None, DataType.BOOL_None) \ .dtype_format(DataType.F16_None, DataType.F16_None, DataType.BOOL_None) \ .dtype_format(DataType.F32_None, DataType.F32_None, DataType.BOOL_None) \ .get_op_info() @op_info_register(less_equal_op_info) def _less_equal_tbe(): """LessEqual TBE register""" return
from te import tik from te import tvm from topi import generic from topi.cce import util # General limitation of the size for input shape: 2**31 SHAPE_SIZE_LIMIT = 2147483648 NoneType = type(None) matmul_cube_dense_left_op_info = TBERegOp("CusMatMulCubeDenseLeft") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("matmulcubedenseleft.so") \ .compute_cost(10) \ .kernel_name("CusMatMulCubeDenseLeft") \ .partial_flag(True) \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .input(2, "x3", False, "optional", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_Default, DataType.F16_FracNZ, DataType.F16_Default, DataType.F16_FracNZ) \ .get_op_info() # pylint: disable=locally-disabled,too-many-arguments,too-many-branches, too-many-statements, too-many-locals, def _shape_check(shape_a, shape_b, shape_bias, src_dtype, trans_a, trans_b): """ Check the given input if legal Parameters: shape_a: list or tuple
# See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ """AsinhGrad op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType asinh_grad_op_info = TBERegOp("AsinhGrad") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("asinh_grad.so") \ .compute_cost(10) \ .kernel_name("asinh_grad") \ .partial_flag(True) \ .input(0, "y", False, "required", "all") \ .input(1, "dy", False, "required", "all") \ .output(0, "z", False, "required", "all") \ .dtype_format(DataType.F16_5HD, DataType.F16_5HD, DataType.F16_5HD) \ .dtype_format(DataType.F16_FracZ, DataType.F16_FracZ, DataType.F16_FracZ) \ .dtype_format(DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0, DataType.F16_C1HWNCoC0) \ .dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ .dtype_format(DataType.F32_5HD, DataType.F32_5HD, DataType.F32_5HD) \ .dtype_format(DataType.F32_FracZ, DataType.F32_FracZ, DataType.F32_FracZ) \ .dtype_format(DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0, DataType.F32_C1HWNCoC0) \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .get_op_info() @op_info_register(asinh_grad_op_info) def _asinh_grad_tbe(): """AsinhGrad TBE register""" return
# 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. # ============================================================================ """batch_matmul_impl""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType from te import tik from topi.cce import util cus_batchmatmul_op_info = TBERegOp("CusBatchMatMul") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("batchmatmul.so") \ .compute_cost(10) \ .kernel_name("CusBatchMatMul") \ .partial_flag(True) \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ .get_op_info() def _get_flattern_shape(shape): """_get_flattern_shape""" flattern_shape = 1 for dim in shape: flattern_shape *= dim return (flattern_shape,)
# 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. # ============================================================================ """ApproximateEqual op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType approximate_equal_op_info = TBERegOp("ApproximateEqual") \ .fusion_type("ELEMWISE") \ .async_flag(False) \ .binfile_name("approximate_equal.so") \ .compute_cost(10) \ .kernel_name("approximate_equal") \ .partial_flag(True) \ .op_pattern("broadcast") \ .attr("tolerance", "optional", "float", "all") \ .input(0, "x1", False, "required", "all") \ .input(1, "x2", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_None, DataType.F16_None, DataType.BOOL_None) \ .dtype_format(DataType.F32_None, DataType.F32_None, DataType.BOOL_None) \ .get_op_info() @op_info_register(approximate_equal_op_info) def _approximate_equal_tbe(): """ApproximateEqual TBE register""" return
# limitations under the License. # ============================================================================ """DepthToSpace op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType depth_to_space_op_info = TBERegOp("DepthToSpace") \ .fusion_type("OPAQUE") \ .async_flag(False) \ .binfile_name("depth_to_space.so") \ .compute_cost(10) \ .kernel_name("depth_to_space") \ .partial_flag(True) \ .attr("block_size", "required", "int", "all") \ .attr("data_format", "optional", "str", "all") \ .input(0, "x", False, "required", "all") \ .output(0, "y", False, "required", "all") \ .dtype_format(DataType.F16_NHWC, DataType.F16_NHWC) \ .dtype_format(DataType.F32_NHWC, DataType.F32_NHWC) \ .dtype_format(DataType.I8_NHWC, DataType.I8_NHWC) \ .dtype_format(DataType.I16_NHWC, DataType.I16_NHWC) \ .dtype_format(DataType.I32_NHWC, DataType.I32_NHWC) \ .dtype_format(DataType.I64_NHWC, DataType.I64_NHWC) \ .dtype_format(DataType.U8_NHWC, DataType.U8_NHWC) \ .dtype_format(DataType.U16_NHWC, DataType.U16_NHWC) \ .dtype_format(DataType.U32_NHWC, DataType.U32_NHWC) \ .dtype_format(DataType.U64_NHWC, DataType.U64_NHWC) \ .get_op_info() @op_info_register(depth_to_space_op_info) def _depth_to_space_tbe(): """DepthToSpace TBE register"""
# 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. # ============================================================================ """Conv3DBackpropInput op""" from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType conv3d_backprop_input_op_info = TBERegOp("Conv3DBackpropInput") \ .fusion_type("CONVLUTION") \ .async_flag(False) \ .binfile_name("conv3d_backprop_input_d.so") \ .compute_cost(10) \ .kernel_name("conv3d_backprop_input_d") \ .partial_flag(True) \ .attr("input_size", "required", "listInt", "all") \ .attr("strides", "required", "listInt", "all") \ .attr("pad_list", "required", "listInt", "all") \ .attr("dilations", "required", "listInt", "all") \ .attr("groups", "optional", "int", "all") \ .attr("format", "optional", "str", "all") \ .input(0, "filter", False, "required", "all") \ .input(1, "out_backprop", False, "required", "all") \ .output(0, "y", True, "required", "all") \ .dtype_format(DataType.F16_FRACTAL_Z_3D, DataType.F16_NDC1HWC0, DataType.F16_NDC1HWC0) \ .get_op_info() @op_info_register(conv3d_backprop_input_op_info) def _conv3d_backprop_input_tbe(): """Conv3DBackpropInput TBE register""" return