Repository: cupy/cupy

Issue:
JIT failed on Jupyter Notebook file.
### Description

Jupyter Notebook output:
```
---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
Input In [1], in <module>
      3 a = cupy.array([0.4, -0.2, 1.8, -1.2])
      4 relu = cupy.vectorize(lambda x: (x > 0.0) * x)
----> 5 print(relu(a))

File ~\miniconda3\envs\cuda\lib\site-packages\cupy\_functional\vectorize.py:87, in vectorize.__call__(self, *args)
     84     raise NotImplementedError
     86 func = _interface._CudaFunction(self.pyfunc, 'numpy', device=True)
---> 87 result = func._emit_code_from_types(in_types, ret_type)
     88 in_params = ', '.join(
     89     f'{t.dtype} in{i}' for i, t in enumerate(in_types))
     90 in_args = ', '.join([f'in{i}' for i in range(len(in_types))])

File ~\miniconda3\envs\cuda\lib\site-packages\cupyx\jit\_interface.py:38, in _CudaFunction._emit_code_from_types(self, in_types, ret_type)
     37 def _emit_code_from_types(self, in_types, ret_type=None):
---> 38     return _compile.transpile(
     39         self.func, self.attributes, self.mode, in_types, ret_type)

File ~\miniconda3\envs\cuda\lib\site-packages\cupyx\jit\_compile.py:123, in transpile(func, attributes, mode, in_types, ret_type)
    121 consts = dict(**cvars.globals, **cvars.nonlocals, **cvars.builtins)
    122 attributes = ' '.join(attributes)
--> 123 tree, source = _parse_function_object(func)
    124 cuda_code, env = _transpile_function(
    125     tree, attributes, mode, consts, in_types, ret_type, source=source)
    126 cuda_code = ''.join([code + '\n' for code in env.preambles]) + cuda_code

File ~\miniconda3\envs\cuda\lib\site-packages\cupyx\jit\_compile.py:88, in _parse_function_object(func)
     85 if filename is None:
     86     raise ValueError(f'JIT needs access to Python source for {func}'
     87                      'but could not be located')
---> 88 with open(filename) as f:
     89     full_source = f.read()
     90 source, start_line = inspect.getsourcelines(func)

FileNotFoundError: [Errno 2] No such file or directory: 'C:\\Users\\honhe\\AppData\\Local\\Temp\\ipykernel_18084\\4273467977.py'
```

### To Reproduce

The source code in Jupyter Notebook file.
```py
import cupy

a = cupy.array([0.4, -0.2, 1.8, -1.2])
relu = cupy.vectorize(lambda x: (x > 0.0) * x)
print(relu(a))  # [ 0.4 -0.   1.8 -0. ]
```


### Installation

Conda-Forge (`conda install ...`)

### Environment

```
> python -c 'import cupy; cupy.show_config()'
C:\Users\honhe\miniconda3\envs\cuda\lib\site-packages\cupy\_environment.py:213: UserWarning: CUDA path could not be detected. Set CUDA_PATH environment variable if CuPy fails to load.
  warnings.warn(
OS                           : Windows-10-10.0.19044-SP0
Python Version               : 3.8.12
CuPy Version                 : 10.0.0
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.22.0
SciPy Version                : None
Cython Build Version         : 0.29.25
Cython Runtime Version       : None
CUDA Root                    : None
nvcc PATH                    : None
CUDA Build Version           : 11020
CUDA Driver Version          : 11020
CUDA Runtime Version         : 11020
cuBLAS Version               : (available)
cuFFT Version                : 10401
cuRAND Version               : 10203
cuSOLVER Version             : (11, 1, 0)
cuSPARSE Version             : (available)
NVRTC Version                : (11, 2)
Thrust Version               : 101000
CUB Build Version            : 101000
Jitify Build Version         : <unknown>
cuDNN Build Version          : None
cuDNN Version                : None
NCCL Build Version           : None
NCCL Runtime Version         : None
cuTENSOR Version             : None
cuSPARSELt Build Version     : None
Device 0 Name                : GeForce MX350
Device 0 Compute Capability  : 61
Device 0 PCI Bus ID          : 0000:01:00.0
```


### Additional Information

_No response_


Discussion:

Resolving patch:
diff --git a/cupyx/jit/__init__.py b/cupyx/jit/__init__.py
--- a/cupyx/jit/__init__.py
+++ b/cupyx/jit/__init__.py
@@ -28,5 +28,4 @@
 from cupyx.jit._builtin_funcs import shfl_down_sync  # NOQA
 from cupyx.jit._builtin_funcs import shfl_xor_sync  # NOQA
 
-_getsource_func = None  # NOQA
 _n_functions_upperlimit = 100
diff --git a/cupyx/jit/_compile.py b/cupyx/jit/_compile.py
--- a/cupyx/jit/_compile.py
+++ b/cupyx/jit/_compile.py
@@ -1,7 +1,7 @@
 import ast
 import collections
 import inspect
-import math
+import linecache
 import numbers
 import re
 import sys
@@ -57,40 +57,61 @@ def new_func(node, *args, **kwargs):
 
 
 def _parse_function_object(func):
-    # Parses function object into ast.FunctionDef object.
+    """Returns the tuple of ``ast.FunctionDef`` object and the source string
+    for the given callable ``func``.
+
+    ``func`` can be a ``def`` function or a ``lambda`` expression.
+
+    The source is returned only for informational purposes (i.e., rendering
+    an exception message in case of an error).
+    """
     if not callable(func):
         raise ValueError('`func` must be a callable object.')
 
+    try:
+        # ``filename`` can be any of:
+        # - A "real" file path on the filesystem
+        # - "<stdin>" (within Python interpreter)
+        # - "<ipython-input-XXXXXXXX>" (within IPython interpreter)
+        filename = inspect.getsourcefile(func)
+    except TypeError:
+        # Built-in function or method, or inside Doctest
+        filename = None
+
+    if filename == '<stdin>':
+        raise RuntimeError(
+            f'JIT needs access to the Python source code for {func}'
+            ' but it cannot be retrieved within the Python interactive'
+            ' interpreter. Consider using IPython instead.')
+
     if func.__name__ != '<lambda>':
-        if jit._getsource_func is None:
-            lines = inspect.getsource(func).split('\n')
-            num_indent = len(lines[0]) - len(lines[0].lstrip())
-            source = '\n'.join([
-                line.replace(' ' * num_indent, '', 1) for line in lines])
-        else:
-            source = jit._getsource_func(func)
+        lines, _ = inspect.getsourcelines(func)
+        num_indent = len(lines[0]) - len(lines[0].lstrip())
+        source = ''.join([
+            line.replace(' ' * num_indent, '', 1) for line in lines])
         tree = ast.parse(source)
         assert isinstance(tree, ast.Module)
         assert len(tree.body) == 1
         return tree.body[0], source
 
-    if jit._getsource_func is not None:
-        full_source = jit._getsource_func(func)
-        start_line, end_line = 0, math.inf
-        source = full_source
-    else:
-        try:
-            filename = inspect.getsourcefile(func)
-        except TypeError:
-            filename = None
-        if filename is None:
-            raise ValueError(f'JIT needs access to Python source for {func}'
-                             'but could not be located')
-        with open(filename) as f:
-            full_source = f.read()
-        source, start_line = inspect.getsourcelines(func)
-        end_line = start_line + len(source)
-        source = ''.join(source)
+    if filename is None:
+        # filename is needed for lambdas.
+        raise ValueError(
+            f'JIT needs access to Python source code for {func}'
+            ' but could not be located.\n'
+            '(hint: it is likely you passed a built-in function or method)')
+
+    # Extract the AST of the lambda from the AST of the whole source file
+    # that defines that lambda.
+    # This is needed because ``inspect.getsourcelines(lambda_expr)`` may
+    # return unparsable code snippet.
+
+    # Use ``linecache.getlines`` instead of directly opening a file to
+    # support notebook environments.
+    full_source = ''.join(linecache.getlines(filename))
+    source, start_line = inspect.getsourcelines(func)
+    end_line = start_line + len(source)
+    source = ''.join(source)
 
     tree = ast.parse(full_source)
 
@@ -98,6 +119,7 @@ def _parse_function_object(func):
              if isinstance(node, ast.Lambda)
              and start_line <= node.lineno < end_line]
     if len(nodes) > 1:
+        # TODO(kmaehashi): can be improved by heuristics (e.g. number of args)
         raise ValueError('Multiple callables are found near the'
                          f' definition of {func}, and JIT could not'
                          ' identify the source code for it.')
@@ -125,7 +147,8 @@ def add_code(self, code: str) -> None:
 
 
 def transpile(func, attributes, mode, in_types, ret_type):
-    """Transpile the target function
+    """Transpiles the target function.
+
     Args:
         func (function): Target function.
         attributes (list of str): Attributes of the generated CUDA function.
@@ -149,10 +172,12 @@ def _transpile_func_obj(func, attributes, mode, in_types, ret_type, generated):
             raise ValueError("Recursive function is not supported.")
         return result
 
+    # Do sanity check first.
+    tree, source = _parse_function_object(func)
+
     cvars = inspect.getclosurevars(func)
     consts = dict(**cvars.globals, **cvars.nonlocals, **cvars.builtins)
     attributes = ' '.join(attributes)
-    tree, source = _parse_function_object(func)
     name = tree.name
     if len(generated.device_function) > 0:
         name += '_' + str(len(generated.device_function))


Test patch: