File tree Expand file tree Collapse file tree 1 file changed +13
-0
lines changed Expand file tree Collapse file tree 1 file changed +13
-0
lines changed Original file line number Diff line number Diff line change 22# Just an early prototype that shows that one can compile elementwise ops into a Metal shader
33from __future__ import annotations
44
5+ import functools
56import itertools
67from typing import Any , Optional , TYPE_CHECKING
78
@@ -701,11 +702,23 @@ def check_bounds(
701702 self .cse .generate (self .compute , line , assignment = False )
702703
703704
705+ @functools .cache
706+ def _warn_prototype () -> None :
707+ import warnings
708+
709+ warnings .warn (
710+ "torch.compile for Metal is an early protoype and might not work as expected."
711+ " For details see https://github.com/pytorch/pytorch/issues/150121" ,
712+ stacklevel = 2 ,
713+ )
714+
715+
704716class MetalScheduling (SIMDScheduling ):
705717 kernel_type = MetalKernel # type: ignore[assignment]
706718
707719 def __init__ (self , scheduler : Optional [Scheduler ]) -> None :
708720 super ().__init__ (scheduler )
721+ _warn_prototype ()
709722 wrapper = V .graph .wrapper_code
710723 if wrapper is not None :
711724 wrapper .header .splice (
You can’t perform that action at this time.
0 commit comments