Prevent example scripts from launching kernels after cute.compile.
()
| 41 | |
| 42 | |
| 43 | def patch_runtime_calls() -> None: |
| 44 | """Prevent example scripts from launching kernels after cute.compile.""" |
| 45 | try: |
| 46 | from cutlass.base_dsl import jit_executor |
| 47 | |
| 48 | def no_run(*_args, **_kwargs): |
| 49 | return None |
| 50 | |
| 51 | jit_executor.JitCompiledFunction.__call__ = no_run |
| 52 | jit_executor.JitExecutor.__call__ = no_run |
| 53 | except Exception: |
| 54 | pass |
| 55 | |
| 56 | try: |
| 57 | from cutlass.cutlass_dsl import tvm_ffi_provider |
| 58 | |
| 59 | def no_run_tvm_ffi(*_args, **_kwargs): |
| 60 | return None |
| 61 | |
| 62 | tvm_ffi_provider.TVMFFIJitCompiledFunction.__call__ = no_run_tvm_ffi |
| 63 | tvm_ffi_provider.TVMFFIJitCompiledFunctionWithKwargs.__call__ = no_run_tvm_ffi |
| 64 | except Exception: |
| 65 | pass |
| 66 | |
| 67 | |
| 68 | def has_compile_option(options: str, option_name: str) -> bool: |