@@ -147,23 +147,22 @@ For the "ahead of time" flavor, we build our C++ extension by writing a
147
147
``setup.py `` script that uses setuptools to compile our C++ code. For the LLTM, it
148
148
looks as simple as this::
149
149
150
- from setuptools import setup
151
- from torch.utils.cpp_extension import CppExtension, BuildExtension
150
+ from setuptools import setup, Extension
151
+ from torch.utils import cpp_extension
152
152
153
153
setup(name='lltm_cpp',
154
- ext_modules=[CppExtension('lltm', ['lltm.cpp'])],
155
- cmdclass={'build_ext': BuildExtension})
156
-
154
+ ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['lltm.cpp'])],
155
+ cmdclass={'build_ext': cpp_extension.BuildExtension})
157
156
158
157
In this code, :class: `CppExtension ` is a convenience wrapper around
159
158
:class: `setuptools.Extension ` that passes the correct include paths and sets
160
159
the language of the extension to C++. The equivalent vanilla :mod: `setuptools `
161
160
code would simply be::
162
161
163
- setuptools. Extension(
162
+ Extension(
164
163
name='lltm_cpp',
165
164
sources=['lltm.cpp'],
166
- include_dirs=torch.utils. cpp_extension.include_paths(),
165
+ include_dirs=cpp_extension.include_paths(),
167
166
language='c++')
168
167
169
168
:class: `BuildExtension ` performs a number of required configuration steps and
@@ -413,7 +412,7 @@ see::
413
412
If we call ``help() `` on the function or module, we can see that its signature
414
413
matches our C++ code::
415
414
416
- In[4] help(lltm .forward)
415
+ In[4] help(lltm_cpp .forward)
417
416
forward(...) method of builtins.PyCapsule instance
418
417
forward(arg0: torch::Tensor, arg1: torch::Tensor, arg2: torch::Tensor, arg3: torch::Tensor, arg4: torch::Tensor) -> List[torch::Tensor]
419
418
@@ -473,6 +472,8 @@ small benchmark to see how much performance we gained from rewriting our op in
473
472
C++. We'll run the LLTM forwards and backwards a few times and measure the
474
473
duration::
475
474
475
+ import time
476
+
476
477
import torch
477
478
478
479
batch_size = 16
0 commit comments