hook up aot compile flow for ats

LiteRT-PiperOrigin-RevId: 819871594

Author: LukeBoyer

litert/ats/BUILD

@@ -39,13 +39,16 @@ cc_test(
     tags = ["no_oss"],
     deps = [
         ":capture",
+        ":compile_fixture",
         ":configure",
         ":inference_fixture",
         ":register",
         "//litert/c:litert_logging",
         "//litert/c:litert_op_code",
+        "//litert/cc:litert_buffer_ref",
         "//litert/cc:litert_c_types_printing",
         "//litert/cc:litert_detail",
+        "//litert/compiler/plugin:compiler_plugin",
         "//litert/test/generators",
         "//litert/test/generators:common",
         "//tflite/schema:schema_fbs",
@@ -67,7 +70,9 @@ cc_library(
         "//litert/cc:litert_expected",
         "//litert/cc:litert_macros",
         "//litert/cc:litert_rng",
+        "//litert/compiler/plugin:compiler_plugin",
         "//litert/core:filesystem_testonly",
+        "//litert/core/model:model_serialize",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/strings",
@@ -101,6 +106,7 @@ cc_library(
     hdrs = ["register.h"],
     deps = [
         ":capture",
+        ":common",
         ":configure",
         "//litert/c:litert_common",
         "//litert/c:litert_logging",
@@ -152,9 +158,26 @@ cc_library(
         "//litert/test:rng_fixture",
         "//litert/test:simple_buffer",
         "//litert/test/generators:common",
-        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
+cc_library(
+    name = "compile_fixture",
+    testonly = True,
+    hdrs = ["compile_fixture.h"],
+    deps = [
+        ":capture",
+        ":common",
+        ":configure",
+        "//litert/c:litert_logging",
+        "//litert/cc:litert_c_types_printing",
+        "//litert/compiler/plugin:compiler_plugin",
+        "//litert/core:filesystem_testonly",
+        "//litert/core/model",
+        "//litert/test:matchers",
+        "//litert/test/generators:common",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -196,19 +219,26 @@ cc_test(
     deps = [
         ":capture",
         ":common",
+        ":compile_fixture",
         ":configure",
+        ":executor",
         ":inference_fixture",
         ":register",
+        "//litert/c:litert_common",
         "//litert/c:litert_op_code",
         "//litert/cc:litert_detail",
         "//litert/cc:litert_expected",
+        "//litert/cc:litert_macros",
+        "//litert/core:filesystem",
+        "//litert/core/model",
+        "//litert/core/model:model_load",
         "//litert/test:common",
+        "//litert/test:simple_buffer",
         "//litert/test/generators",
         "//litert/test/generators:common",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/flags:parse",
-        "@com_google_absl//absl/log:absl_check",
-        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/flags:reflection",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest",
     ],
@@ -239,7 +269,11 @@ cc_library(
     name = "common",
     testonly = True,
     hdrs = ["common.h"],
-    deps = ["@com_google_absl//absl/strings:str_format"],
+    deps = [
+        "//litert/core/model",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
+    ],
 )
 
 cc_library(

litert/ats/ats.cc

@@ -15,25 +15,30 @@
 #include <cstddef>
 #include <cstdint>
 #include <iostream>
+#include <optional>
 
 #include <gtest/gtest.h>
 #include "absl/flags/parse.h"  // from @com_google_absl
 #include "absl/log/absl_check.h"  // from @com_google_absl
 #include "litert/ats/capture.h"
+#include "litert/ats/compile_fixture.h"
 #include "litert/ats/configure.h"
 #include "litert/ats/inference_fixture.h"
 #include "litert/ats/register.h"
 #include "litert/c/litert_logging.h"
 #include "litert/c/litert_op_code.h"
 #include "litert/cc/litert_c_types_printing.h"  // IWYU pragma: keep
 #include "litert/cc/litert_detail.h"
+#include "litert/compiler/plugin/compiler_plugin.h"
 #include "litert/test/generators/common.h"
 #include "litert/test/generators/generators.h"
 #include "tflite/schema/schema_generated.h"
 
 namespace litert::testing {
 namespace {
 
+using ::litert::internal::CompilerPlugin;
+
 static constexpr const char* kArt = R"(
    ###     ######   ######  ######## ##       ######## ########     ###    ########  #######  ########     ######## ########  ######  ########     ######  ##     ## #### ######## ######## 
   ## ##   ##    ## ##    ## ##       ##       ##       ##     ##   ## ##      ##    ##     ## ##     ##       ##    ##       ##    ##    ##       ##    ## ##     ##  ##     ##    ##
@@ -44,23 +49,25 @@ static constexpr const char* kArt = R"(
 ##     ##  ######   ######  ######## ######## ######## ##     ## ##     ##    ##     #######  ##     ##       ##    ########  ######     ##        ######   #######  ####    ##    ######## 
 )";
 
+template <typename Fixture>
 void RegisterNoOp(const AtsConf& options, size_t& test_id, size_t iters,
                   AtsCapture::Ref cap) {
   // clang-format off
   RegisterCombinations<
-      AtsInferenceTest,
+      Fixture,
       NoOp,
       SizeListC<1, 2, 3, 4>,
       TypeList<float, int32_t>>
     (iters, test_id, options, cap);
   // clang-format on
 }
 
+template <typename Fixture>
 void RegisterBinaryNoBroadcast(const AtsConf& options, size_t& test_id,
                                size_t iters, AtsCapture::Ref cap) {
   // clang-format off
   RegisterCombinations<
-      AtsInferenceTest,
+      Fixture,
       BinaryNoBroadcast,
       SizeListC<1, 2, 3, 4, 5, 6>,
       TypeList<float, int32_t>,
@@ -70,6 +77,13 @@ void RegisterBinaryNoBroadcast(const AtsConf& options, size_t& test_id,
   // clang-format on
 }
 
+template <typename Fixture>
+void RegisterAll(const AtsConf& options, size_t& test_id, AtsCapture::Ref cap) {
+  RegisterExtraModels<Fixture>(test_id, options, cap);
+  RegisterNoOp<Fixture>(options, test_id, /*iters=*/10, cap);
+  RegisterBinaryNoBroadcast<Fixture>(options, test_id, /*iters=*/10, cap);
+}
+
 int Ats() {
   std::cerr << kArt << std::endl;
 
@@ -79,9 +93,14 @@ int Ats() {
   size_t test_id = 0;
   AtsCapture cap;
 
-  RegisterNoOp(*options, test_id, /*iters=*/10, cap);
-  RegisterBinaryNoBroadcast(*options, test_id, /*iters=*/10, cap);
-  RegisterExtraModels<AtsInferenceTest>(test_id, *options, cap);
+  std::optional<CompilerPlugin> plugin = std::nullopt;
+
+  if (!options->CompileMode()) {
+    // TODO: lukeboyer - Add compile tests.
+    RegisterAll<AtsInferenceTest>(*options, test_id, cap);
+  } else {
+    RegisterAll<AtsCompileTest>(*options, test_id, cap);
+  }
 
   // Preliminary report.
   {