[JIT] fix resolving of functions in torch/functional. fix compilation of torch.stft (#33504)

Elias Ellison · facebook-github-bot · commit fddf73250d64 · 2020-02-26T18:35:43.000-08:00
Summary: Pull Request resolved: #33504 Fix resolution fo functions that are bound onto torch in torch/functional.py. This does not fix compilation of all of those functions, those will be done in follow ups. Does torch.stft as a start. Fixes #21478 Test Plan: Imported from OSS Differential Revision: D20014591 Pulled By: eellison fbshipit-source-id: bb362f1b5479adbb890e72a54111ef716679d127
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -60,7 +60,7 @@
 from torch.testing._internal.common_utils import run_tests, IS_WINDOWS, TEST_WITH_UBSAN, \
     skipIfRocm, suppress_warnings, IS_SANDCASTLE, GRAPH_EXECUTOR, ProfilingMode, \
     freeze_rng_state, set_rng_seed, slowTest, TemporaryFileName, skipIfCompiledWithoutNumpy, \
-    enable_profiling_mode
+    enable_profiling_mode, TEST_MKL
 from torch.testing._internal.jit_utils import JitTestCase, enable_cpu_fuser, disable_autodiff_subgraph_inlining, \
     _trace, enable_cpu_fuser_if, do_input_map, get_execution_plan, \
     execWrapper, _inline_everything, _tmp_donotuse_dont_inline_everything, \
@@ -10246,6 +10246,15 @@ def test_pack_unpack_state(self):
         self.assertTrue(imported.unpack_called.item())
         torch.testing.assert_allclose(imported(x), x + torch.neg(torch.ones(3, 4, dtype=torch.float)))
 
+    @unittest.skipIf(not TEST_MKL, "PyTorch is built without MKL support")
+    def test_torch_functional(self):
+        def foo(input, n_fft):
+            # type: (Tensor, int) -> Tensor
+            return torch.stft(input, n_fft)
+
+        inps = (torch.randn(10), 7)
+        self.assertEqual(foo(*inps), torch.jit.script(foo)(*inps))
+
     def test_missing_getstate(self):
         class Foo(torch.nn.Module):
             def __init__(self):
diff --git a/torch/_VF.py b/torch/_VF.py
diff --git a/torch/functional.py b/torch/functional.py
@@ -5,6 +5,7 @@
 from ._overrides import has_torch_function, handle_torch_function
 
 Tensor = torch.Tensor
+from torch import _VF
 
 __all__ = [
     'align_tensors',
@@ -301,7 +302,6 @@ def meshgrid(*tensors, **kwargs):
     return torch._C._VariableFunctions.meshgrid(tensors)
 
 
-
 def stft(input, n_fft, hop_length=None, win_length=None, window=None,
          center=True, pad_mode='reflect', normalized=False, onesided=True):
     # type: (Tensor, int, Optional[int], Optional[int], Optional[Tensor], bool, str, bool, bool) -> Tensor
@@ -399,7 +399,7 @@ def stft(input, n_fft, hop_length=None, win_length=None, window=None,
         pad = int(n_fft // 2)
         input = F.pad(input.view(extended_shape), (pad, pad), pad_mode)
         input = input.view(input.shape[-signal_dim:])
-    return torch._C._VariableFunctions.stft(input, n_fft, hop_length, win_length, window, normalized, onesided)
+    return _VF.stft(input, n_fft, hop_length, win_length, window, normalized, onesided)
 
 
 del torch.unique_dim
diff --git a/torch/jit/_builtins.py b/torch/jit/_builtins.py
@@ -83,8 +83,25 @@
     (torch.nn.init._no_grad_zero_, "aten::_no_grad_zero_"),
     (torch._C._get_tracing_state, "aten::_get_tracing_state"),
     (warnings.warn, "aten::warn"),
+    (torch._VF.stft, "aten::stft")
 ]
 
+# ops in torch.functional are bound to torch 
+# in these cases, we want to resolve the function to their python implementation 
+# instead looking up a builtin "aten::" schema
+
+def _gen_torch_functional_registered_ops():
+    # eventually ops should encompass all of torch/functional.py, (torch.functional.__all__) 
+    # but we are currently only able to compile some of the functions. additionally, 
+    # some functions directly map to their aten:: implementations. 
+    # TODO: add support for more ops
+    ops = ["stft"]
+    return set(getattr(torch.functional, name) for name in ops)
+
+_functional_registered_ops = _gen_torch_functional_registered_ops()
+
+def _is_special_functional_bound_op(fn):
+    return fn in _functional_registered_ops
 
 # lazily built to ensure the correct initialization order
 def _get_builtin_table():
@@ -96,7 +113,7 @@ def _get_builtin_table():
     def register_all(mod):
         for name in dir(mod):
             v = getattr(mod, name)
-            if callable(v):
+            if callable(v) and not _is_special_functional_bound_op(v):
                 _builtin_ops.append((v, "aten::" + name))
     for mod in _modules_containing_builtins:
         register_all(mod)
diff --git a/torch/jit/quantized.py b/torch/jit/quantized.py
@@ -2,8 +2,7 @@
 
 from torch._jit_internal import Tuple, Optional, List  # noqa: F401
 
-from torch import Tensor  # noqa: F401
-from torch.nn import _VF
+from torch import Tensor, _VF  # noqa: F401
 
 from torch.nn.utils.rnn import PackedSequence
 
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
@@ -10,7 +10,7 @@
 from .modules import utils
 from .modules.utils import _single, _pair, _triple, _list_with_default
 from . import grad  # noqa: F401
-from . import _VF
+from torch import _VF
 from .._jit_internal import boolean_dispatch, List
 from .._overrides import has_torch_function, handle_torch_function
 
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
@@ -7,7 +7,7 @@
 from ..parameter import Parameter
 from ..utils.rnn import PackedSequence
 from .. import init
-from .. import _VF
+from ... import _VF
 
 _rnn_impls = {
     'RNN_TANH': _VF.rnn_tanh,
diff --git a/torch/nn/quantized/dynamic/modules/rnn.py b/torch/nn/quantized/dynamic/modules/rnn.py
@@ -3,7 +3,7 @@
 import torch
 import torch.nn as nn
 from torch import Tensor  # noqa: F401
-from torch.nn import _VF
+from torch import _VF
 from torch._jit_internal import Tuple, Optional, List  # noqa: F401
 from torch.nn.utils.rnn import PackedSequence
 import numbers
diff --git a/torch/nn/utils/rnn.py b/torch/nn/utils/rnn.py
@@ -2,7 +2,7 @@
 import warnings
 
 import torch
-from .. import _VF
+from ... import _VF
 from ..._jit_internal import Optional
 
 
