Add note and check for GPU capability (#2303)

williamwen42 · malfet · web-flow · commit c26d6a5db19d · 2023-05-01T06:34:08.000-07:00
* add note and check for GPU capability

* Use device capability instead of device name

Co-authored-by: Nikita Shulga &lt;nshulga@meta.com&gt;

* Allow H100

---------

Co-authored-by: Nikita Shulga &lt;nshulga@meta.com&gt;
diff --git a/intermediate_source/torch_compile_tutorial.py b/intermediate_source/torch_compile_tutorial.py
@@ -33,9 +33,25 @@
 # - ``numpy``
 # - ``scipy``
 # - ``tabulate``
-#
-# Note: a modern NVIDIA GPU (Volta or Ampere) is recommended for this tutorial.
-#
+
+######################################################################
+# NOTE: a modern NVIDIA GPU (H100, A100, or V100) is recommended for this tutorial in
+# order to reproduce the speedup numbers shown below and documented elsewhere.
+
+import torch
+import warnings
+
+gpu_ok = False
+if torch.cuda.is_available():
+    device_cap = torch.cuda.get_device_capability()
+    if device_cap in ((7, 0), (8, 0), (9, 0)):
+        gpu_ok = True
+
+if not gpu_ok:
+    warnings.warn(
+        "GPU is not NVIDIA V100, A100, or H100. Speedup numbers may be lower "
+        "than expected."
+    )
 
 ######################################################################
 # Basic Usage
@@ -51,8 +67,6 @@
 # ``torch.compile``. We can then call the returned optimized
 # function in place of the original function.
 
-import torch
-
 def foo(x, y):
     a = torch.sin(x)
     b = torch.cos(x)
