# ============================================================================
"""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