已合并
adapt torch 2.13.0 for catlass in cppwrapper mode #42987
shi-yufeng99创建于 7月27日
adapt torch 2.13.0 for catlass in cppwrapper mode #42987
已合并
共 2 个文件变更+23-9
| @@ -473,6 +473,21 @@ static inline void load_{kernel_name}() {{ | |||
| 473 | self.prefix.writeline("\n") | 473 | self.prefix.writeline("\n") |
| 474 | self.prefix.splice(old_prefix) | 474 | self.prefix.splice(old_prefix) |
| 475 | 475 | ||
| 476 | + def codegen_initialized_kernel_decls(self): | ||
| 477 | + # In AOT mode, CATLASS .o files are linked into model.so, so the | ||
| 478 | + # parent's `extern "C"` declarations are required — same as the | ||
| 479 | + # community CUTLASS flow. | ||
| 480 | + # In JIT cpp_wrapper mode, CATLASS kernels are loaded dynamically | ||
| 481 | + # via dlopen/dlsym through function pointers emitted in | ||
| 482 | + # finalize_prefix(). An `extern "C"` function declaration here | ||
| 483 | + # would conflict with the `static <name>_t <name> = nullptr;` | ||
| 484 | + # pointer variable (C++ forbids a function and a variable sharing | ||
| 485 | + # the same identifier), so we skip it. In JIT mode | ||
| 486 | + # initialized_kernels only contains CATLASS kernels on NPU. | ||
| 487 | + if not V.graph.aot_mode: | ||
| 488 | + return | ||
| 489 | + super().codegen_initialized_kernel_decls() | ||
| 490 | + | ||
| 476 | def codegen_tensor_item( | 491 | def codegen_tensor_item( |
| 477 | self, dtype: torch.dtype, tensor: str, scalar: str, indented_buffer=None | 492 | self, dtype: torch.dtype, tensor: str, scalar: str, indented_buffer=None |
| 478 | ): | 493 | ): |
| @@ -91,15 +91,14 @@ class NPUCombinedScheduling(CUDACombinedScheduling, TritonScheduling): | |||
| 91 | return self._nolinear_triton_scheduling.codegen_node(node) | 91 | return self._nolinear_triton_scheduling.codegen_node(node) |
| 92 | 92 | ||
| 93 | def benchmark_fused_nodes(self, nodes): | 93 | def benchmark_fused_nodes(self, nodes): |
| 94 | - if self.node_can_linear(): | 94 | + try: |
| 95 | - try: | 95 | + return self._triton_scheduling.benchmark_fused_nodes(nodes) |
| 96 | - return self._triton_scheduling.benchmark_fused_nodes(nodes) | 96 | + except Exception: |
| 97 | - except Exception: | 97 | + log.debug( |
| 98 | - log.debug( | 98 | + "linear benchmark_fused_nodes for nodes %s raise error, fallback to origin benchmark_fused_nodes", |
| 99 | - "linear benchmark_fused_nodes for nodes %s raise error, fallback to origin benchmark_fused_nodes", | 99 | + {n.get_name() for n in nodes}, |
| 100 | - {n.get_name() for n in nodes}, | 100 | + exc_info=True, |
| 101 | - exc_info=True, | 101 | + ) |
| 102 | - ) | ||
| 103 | # regroup snode | 102 | # regroup snode |
| 104 | for node in nodes: | 103 | for node in nodes: |
| 105 | for snode in node.get_nodes(): | 104 | for snode in node.get_nodes(): |