Update

[ghstack-poisoned]

Author: swolchok

.ci/scripts/unittest-buck2.sh

@@ -12,9 +12,9 @@ set -eux
 buck2 query "//backends/apple/... + //backends/example/... + \
 //backends/mediatek/... + //backends/transforms/... + \
 //backends/xnnpack/... + //configurations/... + //extension/flat_tensor: + \
-//kernels/aten/... + //kernels/optimized/... + //kernels/portable/... + \
-//kernels/quantized/... + //kernels/test/... + //runtime/... + //schema/... \
-+ //test/... + //util/..."
+//extension/llm/runner: + //kernels/aten/... + //kernels/optimized/... + \
+//kernels/portable/... + //kernels/quantized/... + //kernels/test/... + \
+//runtime/... + //schema/... + //test/... + //util/..."
 
 # TODO: optimized ops are unbuildable because they now use ATen; put
 # them back after we can use PyTorch in OSS buck.

extension/llm/runner/io_manager/targets.bzl

@@ -1,8 +1,8 @@
-load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime")
 
 def define_common_targets():
 
-    for aten in (True, False):
+    for aten in get_aten_mode_options():
         aten_suffix = "_aten" if aten else ""
 
         # Interface for IOManager. No concrete impl from this dep.

extension/llm/runner/targets.bzl

@@ -1,4 +1,4 @@
-load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime")
 
 def define_common_targets():
     runtime.cxx_library(
@@ -32,7 +32,7 @@ def define_common_targets():
         ],
     )
 
-    for aten in (True, False):
+    for aten in get_aten_mode_options():
         aten_suffix = "_aten" if aten else ""
 
         runtime.cxx_library(