Add AMD GPU model and gfx code documentation

Add a comment documenting supported AMD GPU models and their corresponding LLVM gfx codes, including Navi4, MI300X, and MI350.
pytorch · Feb 25, 2025 · d1e92cf · d1e92cf
1 parent 438f2c5
commit d1e92cf
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/torchao/utils.py b/torchao/utils.py
@@ -603,6 +603,13 @@ def fill_defaults(args, n, defaults_tail):
 def _torch_version_at_least(min_version):
     return is_fbcode() or version("torch") >= min_version
 
+# Supported AMD GPU Models and their LLVM gfx Codes:
+#
+# | AMD GPU Model | LLVM gfx Code          |
+# |---------------|------------------------|
+# | Navi4         | gfx1200, gfx1201       |
+# | MI300X        | gfx940, gfx941, gfx942 |
+# | MI350         | gfx950                 |
 
 def is_MI300():
     if torch.cuda.is_available() and torch.version.hip: