已合并
fix add COMPATIBLE for autograd #34359
huangyunlong创建于 4月25日
fix add COMPATIBLE for autograd #34359
已合并
共 1 个文件变更+13-1
| @@ -42,6 +42,14 @@ ${return_type} ${type_wrapper_name}(${formals}); | |||
| 42 | """) | 42 | """) |
| 43 | 43 | ||
| 44 | 44 | ||
| 45 | +OVERWRITE_API_LIST = [ | ||
| 46 | + 'matmul', | ||
| 47 | + 'matmul.out', | ||
| 48 | + 'matmul_backward', | ||
| 49 | + 'matmul_double_backward', | ||
| 50 | +] | ||
| 51 | + | ||
| 52 | + | ||
| 45 | def gen_variable_type( | 53 | def gen_variable_type( |
| 46 | out: str, | 54 | out: str, |
| 47 | fns_with_diff_infos: List[NativeFunctionWithDifferentiabilityInfo], | 55 | fns_with_diff_infos: List[NativeFunctionWithDifferentiabilityInfo], |
| @@ -125,7 +133,11 @@ def gen_variable_type_func( | |||
| 125 | type_definition = type_definition.replace('at::redispatch', 'at_npu::redispatch') | 133 | type_definition = type_definition.replace('at::redispatch', 'at_npu::redispatch') |
| 126 | type_definition = type_definition.replace('_symint', '') | 134 | type_definition = type_definition.replace('_symint', '') |
| 127 | 135 | ||
| 128 | - wrapper_registration = gen_wrapper_registration(f, "Default") | 136 | + if str(f.func.name) in OVERWRITE_API_LIST: |
| 137 | + wrapper_registration = f"""if (std::getenv("TORCH_NPU_USE_COMPATIBLE_IMPL") == nullptr || std::string(std::getenv("TORCH_NPU_USE_COMPATIBLE_IMPL")) != "1") {{ | ||
| 138 | + {gen_wrapper_registration(f, "Default")}}}""" | ||
| 139 | + else: | ||
| 140 | + wrapper_registration = gen_wrapper_registration(f, "Default") | ||
| 129 | result[f"type_derived_method_definitions"] = [type_definition] | 141 | result[f"type_derived_method_definitions"] = [type_definition] |
| 130 | if str(f.func.name) in NPU_AUTOGRAD_FUNCTION: | 142 | if str(f.func.name) in NPU_AUTOGRAD_FUNCTION: |
| 131 | result[f"wrapper_registrations_aten"] = [] | 143 | result[f"wrapper_registrations_aten"] = [] |