[Target] Add Jetson AGX Orin tags (#16231)

apache · Dec 15, 2023 · 37c38ac · 37c38ac
1 parent c8bfdb2
commit 37c38ac
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/src/target/tag.cc b/src/target/tag.cc
@@ -106,6 +106,30 @@ TVM_REGISTER_TARGET_TAG("nvidia/jetson-orin-nano")
                                                  {"mcpu", String("carmel")},
                                                  {"num-cores", Integer(6)}}}});
 
+TVM_REGISTER_TARGET_TAG("nvidia/jetson-agx-orin-32gb")
+    .set_config({{"kind", String("cuda")},
+                 {"arch", String("sm_87")},
+                 {"max_shared_memory_per_block", Integer(49152)},
+                 {"max_threads_per_block", Integer(1024)},
+                 {"thread_warp_size", Integer(32)},
+                 {"registers_per_block", Integer(65536)},
+                 {"host", Map<String, ObjectRef>{{"kind", String("llvm")},
+                                                 {"mtriple", String("aarch64-linux-gnu")},
+                                                 {"mcpu", String("cortex-a78")},
+                                                 {"num-cores", Integer(8)}}}});
+
+TVM_REGISTER_TARGET_TAG("nvidia/jetson-agx-orin-64gb")
+    .set_config({{"kind", String("cuda")},
+                 {"arch", String("sm_87")},
+                 {"max_shared_memory_per_block", Integer(49152)},
+                 {"max_threads_per_block", Integer(1024)},
+                 {"thread_warp_size", Integer(32)},
+                 {"registers_per_block", Integer(65536)},
+                 {"host", Map<String, ObjectRef>{{"kind", String("llvm")},
+                                                 {"mtriple", String("aarch64-linux-gnu")},
+                                                 {"mcpu", String("cortex-a78")},
+                                                 {"num-cores", Integer(12)}}}});
+
 #define TVM_REGISTER_CUDA_TAG(Name, Arch, SharedMem, RegPerBlock) \
   TVM_REGISTER_TARGET_TAG(Name).set_config({                      \
       {"kind", String("cuda")},                                   \