^jddlZddlZddlZddlmZddlmZddlmZ ddl m Z m Z m Z mZddlmZmZdd lmZmZmZmZej0eZd Zd Zd Zej<d ee edzffdZ!d e fdZ"d dezfdZ#ej<d e fdZ$ddZdddde%de%de%dedzdede de d dfdZejLd eeejNey)N)cast)Version)cuda)_available_version_unavailable_reasoncheck_native_jit_disabledcheck_native_version_skip) dsl_registryDSLModuleProtocol) _OpCondFn _OpImplFnderegister_op_overridesregister_op_overridetritonreturnctjsydg}t|}|d}td}||fStj d|d}d}||fS)zO Check if triton is available NOTE: must not import at this point )FN)rrNTrz*triton native DSL ops require: `triton` %sF)_cudais_builtr rloginfo)depsreason availableversions e/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/torch/_native/triton_utils.py_check_runtime_availabler sp >>   D! &F ~ $X. g  =vF  g c t\}}|SNr )r_s rruntime_availabler&<s+-LIq r!c t\}}|Sr#r$)r%rs rruntime_versionr(As)+JAw Nr!ct\}}|y|jtk(}|jtk\}|r|s t ryt jd|tty)NFTzetriton version %s is not sufficient (>= (%s.%s.*)); set TORCH_NATIVE_SKIP_VERSION_CHECK=1 to override)r major_TRITON_REQUIRED_VERSION_MAJORminor_TRITON_MINIMUM_VERSION_MINORr rr)r%rmajor_okminor_oks r_version_is_sufficientr0Fsc)+JAw}} >>H}} ==HX";"=HH <&%  r!c$tty)z+ Deregister all ops through triton )disable_dsl_namesN)_deregister_op_overrides_impl_TRITON_DSL_NAMEr!rrr^s "4DEr!Fallow_multiple_overrideunconditional_override lib_symbol op_symbol dispatch_keycondimplr7r8c |t\}}|r trytsytt|||||||y)z See torch/_native/registry.py for the underlying implementation and arguments. This is a thin, DSL-checking wrapper over _register_op_override_impl Nr6)r r r0_register_op_override_implr4) r9r:r;r<r=r7r8rrs rrresF23Iw 35 ! #  75 r!)rN)( functoolsloggingsystypingrtorch._vendor.packaging.versionrbackendsrr common_utilsrr r r r r registryrrrr3rr? getLogger__name__rr4r+r-cachetupleboolr r&r(r0str register_dslmodulesr5r!rrPsJ 3$ :g!!" ! %gn(<"=64   .F%*#( d    "! H (D):CKK