mindspore 2.3.0__cp39-none-any.whl → 2.3.0rc2__cp39-none-any.whl

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl/dynamic/decoder_kv_cache.py

@@ -37,37 +37,18 @@ DTYPE_MAP = {"float32": ["DT_FLOAT", "float"],
     "quint16": ["DT_QUINT16", "unknown"],
     "resource": ["DT_RESOURCE", "unknown"],
     "string_ref": ["DT_STRING_REF", "unknown"],
-    "int4": ["DT_INT4", "int4b_t"],
+    "int4": ["DT_INT4", "int8_t"],
     "bfloat16": ["DT_BF16", "bfloat16_t"]}
 
-def add_dtype_fmt_option_single(x, x_n, is_ref: bool = False):
-    options = []
-    x_fmt = x.get("format")
-    x_dtype = x.get("dtype")
-    x_n_in_kernel = x_n + '_REF' if is_ref else x_n
-    options.append("-DDTYPE_{n}={t}".format(n=x_n_in_kernel, t=DTYPE_MAP.get(x_dtype)[1]))
-    options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n_in_kernel, ot=DTYPE_MAP.get(x_dtype)[0]))
-    options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n_in_kernel, f=x_fmt))
-    return options
- 
 def get_dtype_fmt_options(__inputs__, __outputs__):
     options = []
-    unique_param_name_set = set()
-    for x in __inputs__:
-        if x is None:
-            continue
-        x_n = x.get("param_name")[:-5].upper()
-        unique_param_name_set.add(x_n)
-        options += add_dtype_fmt_option_single(x, x_n)
- 
-    for x in __outputs__:
-        if x is None:
-            continue
-        x_n = x.get("param_name")[:-5].upper()
-        if x_n in unique_param_name_set:
-            options += add_dtype_fmt_option_single(x, x_n, True)
-        else:
-            options += add_dtype_fmt_option_single(x, x_n)
+    for x in __inputs__ + __outputs__:
+        x_n = x.get("param_name").upper()
+        x_fmt = x.get("format")
+        x_dtype = x.get("dtype")
+        options.append("-DDTYPE_{n}={t}".format(n=x_n, t=DTYPE_MAP.get(x_dtype)[1]))
+        options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n, ot=DTYPE_MAP.get(x_dtype)[0]))
+        options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n, f=x_fmt))
     return options
 
 def load_dso(so_path):
@@ -79,40 +60,20 @@ def load_dso(so_path):
     else:
         print("load so succ ", so_path)
 
-def get_shortsoc_compile_option(compile_option_list: list, shortsoc:str):
-    compile_options = []
-    if shortsoc in compile_option_list:
-        compile_options = compile_option_list[shortsoc]
-    elif '__ALLSOC__' in compile_option_list:
-        compile_options = compile_option_list['__ALLSOC__']
-    return compile_options
-
-def get_kernel_source(src_file, dir_snake, dir_ex):
-    src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, src_file)
-    if os.path.exists(src_ex):
-        return src_ex
-    src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, src_file)
-    if os.path.exists(src):
-        return src
-    src = os.path.join(PYF_PATH, src_file)
-    if os.path.exists(src):
-        return src
-    return src_ex
-
-def _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_):
+def _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out):
     __inputs__ = []
-    for arg in [cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__]:
+    for arg in [cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len]:
         if arg != None:
-            if isinstance(arg, (list, tuple)):
+            if type(arg) is list:
                 if len(arg) == 0:
                     continue
                 __inputs__.append(arg[0])
             else:
                 __inputs__.append(arg)
     __outputs__ = []
-    for arg in [out_out_]:
+    for arg in [out]:
         if arg != None:
-            if isinstance(arg, (list, tuple)):
+            if type(arg) is list:
                 if len(arg) == 0:
                     continue
                 __outputs__.append(arg[0])
@@ -121,18 +82,20 @@ def _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, s
     __attrs__ = []
     return __inputs__, __outputs__, __attrs__
 
-@tbe_register.register_operator("DecoderKvCache", trans_bool_to_s8=False)
+@tbe_register.register_operator("DecoderKvCache")
 @para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
-def decoder_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, kernel_name="decoder_kv_cache", impl_mode=""):
+def decoder_kv_cache(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, kernel_name="decoder_kv_cache", impl_mode=""):
     if get_current_build_config("enable_op_prebuild"):
         return
-    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
     options = get_dtype_fmt_options(__inputs__, __outputs__)
     options += ["-x", "cce"]
-    bisheng = shutil.which("bisheng")
-    if bisheng != None:
-        bisheng_path = os.path.dirname(bisheng)
-        tikcpp_path = os.path.realpath(os.path.join(bisheng_path, "..", "..", "tikcpp"))
+    ccec = os.environ.get('CCEC_REAL_PATH')
+    if ccec is None:
+        ccec = shutil.which("ccec")
+    if ccec != None:
+        ccec_path = os.path.dirname(ccec)
+        tikcpp_path = os.path.realpath(os.path.join(ccec_path, "..", "..", "tikcpp"))
     else:
         tikcpp_path = os.path.realpath("/usr/local/Ascend/latest/compiler/tikcpp")
     options.append("-I" + tikcpp_path)
@@ -146,35 +109,26 @@ def decoder_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in
         options.append("-DHIGH_PRECISION=1")
     if get_default_build_config("enable_deterministic_mode") == 1:
         options.append("-DDETEMINISTIC_MODE=1")
-
-    custom_compile_options = {},
-    custom_all_compile_options = {},
-    soc_version = get_soc_spec("SOC_VERSION")
-    soc_short = get_soc_spec("SHORT_SOC_VERSION").lower()
-    custom_compile_options_soc = get_shortsoc_compile_option(custom_compile_options[0], soc_short)
-    custom_all_compile_options_soc = get_shortsoc_compile_option(custom_all_compile_options[0], soc_short)
-    options += custom_all_compile_options_soc
-    options += custom_compile_options_soc
-
     origin_func_name = "decoder_kv_cache"
-    ascendc_src_dir_ex = "decoder_kv_cache"
     ascendc_src_dir = "decoder_kv_cache"
     ascendc_src_file = "decoder_kv_cache.cpp"
-    src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex)
+    src = os.path.join(PYF_PATH, "..", "ascendc", ascendc_src_dir, ascendc_src_file)
+    if not os.path.exists(src):
+        src = os.path.join(PYF_PATH, ascendc_src_file)
 
-    print("start compile Ascend C operator DecoderKvCache. kernel name is " + kernel_name)
+    print("start compile Ascend C operator DecoderKvCache. kernel name is decoder_kv_cache")
     op_type = "DecoderKvCache"
     code_channel = get_code_channel(src, kernel_name, op_type, options)
     op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\
-        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__], origin_outputs = [out_out_])
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len], origin_outputs = [out])
     compile_op(src, origin_func_name, op_info, options, code_channel, '{}')
 
-def op_select_format(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):
-    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+def op_select_format(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
     result = check_op_cap("op_select_format", "DecoderKvCache", __inputs__, __outputs__, __attrs__)
     return result.decode("utf-8")
 
-def get_op_specific_info(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):
-    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+def get_op_specific_info(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
     result = check_op_cap("get_op_specific_info", "DecoderKvCache", __inputs__, __outputs__, __attrs__)
     return result.decode("utf-8")

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl/dynamic/prompt_kv_cache.py

@@ -37,37 +37,18 @@ DTYPE_MAP = {"float32": ["DT_FLOAT", "float"],
     "quint16": ["DT_QUINT16", "unknown"],
     "resource": ["DT_RESOURCE", "unknown"],
     "string_ref": ["DT_STRING_REF", "unknown"],
-    "int4": ["DT_INT4", "int4b_t"],
+    "int4": ["DT_INT4", "int8_t"],
     "bfloat16": ["DT_BF16", "bfloat16_t"]}
 
-def add_dtype_fmt_option_single(x, x_n, is_ref: bool = False):
-    options = []
-    x_fmt = x.get("format")
-    x_dtype = x.get("dtype")
-    x_n_in_kernel = x_n + '_REF' if is_ref else x_n
-    options.append("-DDTYPE_{n}={t}".format(n=x_n_in_kernel, t=DTYPE_MAP.get(x_dtype)[1]))
-    options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n_in_kernel, ot=DTYPE_MAP.get(x_dtype)[0]))
-    options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n_in_kernel, f=x_fmt))
-    return options
- 
 def get_dtype_fmt_options(__inputs__, __outputs__):
     options = []
-    unique_param_name_set = set()
-    for x in __inputs__:
-        if x is None:
-            continue
-        x_n = x.get("param_name")[:-5].upper()
-        unique_param_name_set.add(x_n)
-        options += add_dtype_fmt_option_single(x, x_n)
- 
-    for x in __outputs__:
-        if x is None:
-            continue
-        x_n = x.get("param_name")[:-5].upper()
-        if x_n in unique_param_name_set:
-            options += add_dtype_fmt_option_single(x, x_n, True)
-        else:
-            options += add_dtype_fmt_option_single(x, x_n)
+    for x in __inputs__ + __outputs__:
+        x_n = x.get("param_name").upper()
+        x_fmt = x.get("format")
+        x_dtype = x.get("dtype")
+        options.append("-DDTYPE_{n}={t}".format(n=x_n, t=DTYPE_MAP.get(x_dtype)[1]))
+        options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n, ot=DTYPE_MAP.get(x_dtype)[0]))
+        options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n, f=x_fmt))
     return options
 
 def load_dso(so_path):
@@ -79,40 +60,20 @@ def load_dso(so_path):
     else:
         print("load so succ ", so_path)
 
-def get_shortsoc_compile_option(compile_option_list: list, shortsoc:str):
-    compile_options = []
-    if shortsoc in compile_option_list:
-        compile_options = compile_option_list[shortsoc]
-    elif '__ALLSOC__' in compile_option_list:
-        compile_options = compile_option_list['__ALLSOC__']
-    return compile_options
-
-def get_kernel_source(src_file, dir_snake, dir_ex):
-    src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, src_file)
-    if os.path.exists(src_ex):
-        return src_ex
-    src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, src_file)
-    if os.path.exists(src):
-        return src
-    src = os.path.join(PYF_PATH, src_file)
-    if os.path.exists(src):
-        return src
-    return src_ex
-
-def _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_):
+def _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out):
     __inputs__ = []
-    for arg in [cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__]:
+    for arg in [cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len]:
         if arg != None:
-            if isinstance(arg, (list, tuple)):
+            if type(arg) is list:
                 if len(arg) == 0:
                     continue
                 __inputs__.append(arg[0])
             else:
                 __inputs__.append(arg)
     __outputs__ = []
-    for arg in [out_out_]:
+    for arg in [out]:
         if arg != None:
-            if isinstance(arg, (list, tuple)):
+            if type(arg) is list:
                 if len(arg) == 0:
                     continue
                 __outputs__.append(arg[0])
@@ -121,18 +82,20 @@ def _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, s
     __attrs__ = []
     return __inputs__, __outputs__, __attrs__
 
-@tbe_register.register_operator("PromptKvCache", trans_bool_to_s8=False)
+@tbe_register.register_operator("PromptKvCache")
 @para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
-def prompt_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, kernel_name="prompt_kv_cache", impl_mode=""):
+def prompt_kv_cache(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, kernel_name="prompt_kv_cache", impl_mode=""):
     if get_current_build_config("enable_op_prebuild"):
         return
-    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
     options = get_dtype_fmt_options(__inputs__, __outputs__)
     options += ["-x", "cce"]
-    bisheng = shutil.which("bisheng")
-    if bisheng != None:
-        bisheng_path = os.path.dirname(bisheng)
-        tikcpp_path = os.path.realpath(os.path.join(bisheng_path, "..", "..", "tikcpp"))
+    ccec = os.environ.get('CCEC_REAL_PATH')
+    if ccec is None:
+        ccec = shutil.which("ccec")
+    if ccec != None:
+        ccec_path = os.path.dirname(ccec)
+        tikcpp_path = os.path.realpath(os.path.join(ccec_path, "..", "..", "tikcpp"))
     else:
         tikcpp_path = os.path.realpath("/usr/local/Ascend/latest/compiler/tikcpp")
     options.append("-I" + tikcpp_path)
@@ -146,35 +109,26 @@ def prompt_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in_
         options.append("-DHIGH_PRECISION=1")
     if get_default_build_config("enable_deterministic_mode") == 1:
         options.append("-DDETEMINISTIC_MODE=1")
-
-    custom_compile_options = {},
-    custom_all_compile_options = {},
-    soc_version = get_soc_spec("SOC_VERSION")
-    soc_short = get_soc_spec("SHORT_SOC_VERSION").lower()
-    custom_compile_options_soc = get_shortsoc_compile_option(custom_compile_options[0], soc_short)
-    custom_all_compile_options_soc = get_shortsoc_compile_option(custom_all_compile_options[0], soc_short)
-    options += custom_all_compile_options_soc
-    options += custom_compile_options_soc
-
     origin_func_name = "prompt_kv_cache"
-    ascendc_src_dir_ex = "prompt_kv_cache"
     ascendc_src_dir = "prompt_kv_cache"
     ascendc_src_file = "prompt_kv_cache.cpp"
-    src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex)
+    src = os.path.join(PYF_PATH, "..", "ascendc", ascendc_src_dir, ascendc_src_file)
+    if not os.path.exists(src):
+        src = os.path.join(PYF_PATH, ascendc_src_file)
 
-    print("start compile Ascend C operator PromptKvCache. kernel name is " + kernel_name)
+    print("start compile Ascend C operator PromptKvCache. kernel name is prompt_kv_cache")
     op_type = "PromptKvCache"
     code_channel = get_code_channel(src, kernel_name, op_type, options)
     op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\
-        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__], origin_outputs = [out_out_])
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len], origin_outputs = [out])
     compile_op(src, origin_func_name, op_info, options, code_channel, '{}')
 
-def op_select_format(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):
-    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+def op_select_format(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
     result = check_op_cap("op_select_format", "PromptKvCache", __inputs__, __outputs__, __attrs__)
     return result.decode("utf-8")
 
-def get_op_specific_info(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):
-    __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_)
+def get_op_specific_info(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
     result = check_op_cap("get_op_specific_info", "PromptKvCache", __inputs__, __outputs__, __attrs__)
     return result.decode("utf-8")

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_proto/inc/op_proto.h

@@ -6,10 +6,11 @@
 
 namespace ge {
 
-REG_OP(AllFinite)
-    .INPUT(gradient, ge::TensorType::ALL())
-    .OUTPUT(is_finite, ge::TensorType::ALL())
-    .OP_END_FACTORY_REG(AllFinite);
+REG_OP(AddCustom)
+    .INPUT(x, ge::TensorType::ALL())
+    .INPUT(y, ge::TensorType::ALL())
+    .OUTPUT(z, ge::TensorType::ALL())
+    .OP_END_FACTORY_REG(AddCustom);
 
 REG_OP(DecoderKvCache)
     .INPUT(cache, ge::TensorType::ALL())