Skip to content

Commit 731ee52

Browse files
terryheocopybara-github
terryheo
authored and
copybara-github
committed
No default CPU Acceleration
Xnnpack was applied by default since it didn't report 'IsTfLiteDelegateResponsibleForJitCompilation' correctly. Now it's only applied when CPU Accelerator is specified. Added a test to check the error when model can't be fully accelerated. LiteRT-PiperOrigin-RevId: 828256860
1 parent 25c53d3 commit 731ee52

File tree

6 files changed

+153
-49
lines changed

6 files changed

+153
-49
lines changed

‎litert/cc/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,7 @@ litert_device_test(
357357
# "requires-gpu-nvidia",
358358
# ],
359359
# deps = [
360+
# ":litert_common",
360361
# ":litert_compiled_model",
361362
# ":litert_element_type",
362363
# ":litert_environment",

‎litert/cc/litert_compiled_model_gpu_test.cc‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "litert/c/litert_event_type.h"
3030
#include "litert/c/litert_profiler_event.h"
3131
#include "litert/cc/internal/litert_platform_support.h"
32+
#include "litert/cc/litert_common.h"
3233
#include "litert/cc/litert_compiled_model.h"
3334
#include "litert/cc/litert_element_type.h"
3435
#include "litert/cc/litert_environment.h"
@@ -386,6 +387,28 @@ TEST_P(CompiledModelGpuTest, PartialDelegation) {
386387
}
387388
}
388389

390+
TEST_P(CompiledModelGpuTest, PartialDelegationNoCpuFallbackError) {
391+
constexpr const char* kModelPartilaFileName = "simple_cast_and_add_op.tflite";
392+
LITERT_ASSERT_OK_AND_ASSIGN(
393+
auto model,
394+
Model::CreateFromFile(testing::GetTestFilePath(kModelPartilaFileName)));
395+
396+
auto env = litert::Environment::Create({});
397+
ASSERT_TRUE(env);
398+
399+
auto compilation_options = Options::Create();
400+
compilation_options->SetHardwareAccelerators(HwAccelerators::kGpu);
401+
LITERT_ASSERT_OK_AND_ASSIGN(auto gpu_options, litert::GpuOptions::Create());
402+
LITERT_ASSERT_OK(
403+
gpu_options.EnableExternalTensorsMode(CompiledModelGpuTest::GetParam()));
404+
compilation_options->AddOpaqueOptions(std::move(gpu_options));
405+
406+
auto compiled_model_res =
407+
CompiledModel::Create(*env, model, *compilation_options);
408+
EXPECT_FALSE(compiled_model_res.HasValue());
409+
EXPECT_EQ(compiled_model_res.Error().Status(), kLiteRtStatusErrorCompilation);
410+
}
411+
389412
TEST_P(CompiledModelGpuTest, BasicAdd3dCstInt32) {
390413
constexpr const char* kInt32ModelFileName = "simple_add3d_cst_int32.tflite";
391414
constexpr const int32_t kInt32TestInput0Tensor[] = {1, 2, 3, 4, 5, 6};

‎litert/runtime/accelerators/xnnpack/xnnpack_accelerator.cc‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,17 @@ class CpuAccelerator final
109109
LiteRtUnwrapDelegate(delegate_wrapper, &xnnpack_delegate);
110110
TfLiteXNNPackDelegateDelete(xnnpack_delegate);
111111
}
112+
113+
// Returns true to indicate the XNNPack delegate is responsible for JIT
114+
// compilation.
115+
static LiteRtStatus IsTfLiteDelegateResponsibleForJitCompilation(
116+
LiteRtAcceleratorT* accelerator, bool* does_jit_compilation) {
117+
LITERT_RETURN_IF_ERROR(does_jit_compilation,
118+
litert::ErrorStatusBuilder::InvalidArgument())
119+
<< "`does_jit_compilation` pointer is null.";
120+
*does_jit_compilation = true;
121+
return kLiteRtStatusOk;
122+
}
112123
};
113124

114125
} // namespace
@@ -131,6 +142,11 @@ LiteRtStatus LiteRtRegisterCpuAccelerator(LiteRtEnvironment environment) {
131142
LITERT_ASSIGN_OR_RETURN(auto accelerator_impl,
132143
litert::CpuAccelerator::Create());
133144

145+
LITERT_RETURN_IF_ERROR(
146+
LiteRtSetIsAcceleratorDelegateResponsibleForJitCompilation(
147+
accelerator.get(), litert::CpuAccelerator::
148+
IsTfLiteDelegateResponsibleForJitCompilation));
149+
134150
LITERT_RETURN_IF_ERROR(LiteRtRegisterAccelerator(
135151
environment, accelerator.release(), accelerator_impl.release(),
136152
litert::CpuAccelerator::Destroy));

‎litert/tools/BUILD‎

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -399,25 +399,35 @@ cc_library(
399399
# )
400400
# copybara:uncomment_end
401401

402-
#copybara:comment_begin(google-only)
403402
cc_test(
404403
name = "benchmark_litert_model_test",
405404
srcs = ["benchmark_litert_model_test.cc"],
406405
data = ["//litert/test:testdata/mobilenet_v2_1.0_224.tflite"],
407-
tags = ["requires-gpu-nvidia"],
408-
deps =
409-
[
410-
":benchmark_litert_model",
411-
"@com_google_googletest//:gtest_main",
412-
# copybara:uncomment_begin(google-only)
413-
# "//litert/runtime/accelerators/gpu:ml_drift_cl_accelerator", # buildcleaner: keep
414-
# copybara:uncomment_end
415-
"//tflite/core/c:private_c_api_types",
416-
"//tflite/tools/benchmark:benchmark_model_lib",
417-
"//tflite/tools/benchmark:benchmark_params",
418-
],
406+
deps = [
407+
":benchmark_litert_model",
408+
"//tflite/core/c:private_c_api_types",
409+
"//tflite/tools/benchmark:benchmark_model_lib",
410+
"//tflite/tools/benchmark:benchmark_params",
411+
"@com_google_googletest//:gtest_main",
412+
],
419413
)
420-
#copybara:comment_end
414+
415+
# copybara:uncomment_begin(google-only)
416+
# cc_test(
417+
# name = "benchmark_litert_model_gpu_test",
418+
# srcs = ["benchmark_litert_model_gpu_test.cc"],
419+
# data = ["//litert/test:testdata/mobilenet_v2_1.0_224.tflite"],
420+
# tags = ["requires-gpu-nvidia"],
421+
# deps = [
422+
# ":benchmark_litert_model",
423+
# "@com_google_googletest//:gtest_main",
424+
# "//litert/runtime/accelerators/gpu:ml_drift_cl_accelerator", # buildcleaner: keep
425+
# "//tflite/core/c:private_c_api_types",
426+
# "//tflite/tools/benchmark:benchmark_model_lib",
427+
# "//tflite/tools/benchmark:benchmark_params",
428+
# ],
429+
# )
430+
# copybara:uncomment_end
421431

422432
cc_library(
423433
name = "tensor_utils",

‎litert/tools/benchmark_litert_model_gpu_test.cc‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "litert/tools/benchmark_litert_model.h"
17+
18+
#include <fcntl.h>
19+
#include <sys/stat.h>
20+
21+
22+
#include <string>
23+
#include <utility>
24+
25+
#include <gtest/gtest.h>
26+
#include "tflite/core/c/c_api_types.h"
27+
#include "tflite/tools/benchmark/benchmark_model.h"
28+
#include "tflite/tools/benchmark/benchmark_params.h"
29+
30+
namespace litert {
31+
namespace benchmark {
32+
namespace {
33+
using ::litert::benchmark::BenchmarkLiteRtModel;
34+
using ::tflite::benchmark::BenchmarkListener;
35+
using ::tflite::benchmark::BenchmarkParams;
36+
using ::tflite::benchmark::BenchmarkResults;
37+
38+
static constexpr char kModelPath[] =
39+
"third_party/odml/litert/litert/test/testdata/"
40+
"mobilenet_v2_1.0_224.tflite";
41+
static constexpr char kSignatureToRunFor[] = "<placeholder signature>";
42+
43+
class TestBenchmarkListener : public BenchmarkListener {
44+
public:
45+
void OnBenchmarkEnd(const BenchmarkResults& results) override {
46+
results_ = results;
47+
}
48+
49+
BenchmarkResults results_;
50+
};
51+
52+
TEST(BenchmarkLiteRtModelTest, GPUAcceleration) {
53+
// MSAN does not support GPU tests.
54+
#if defined(MEMORY_SANITIZER) || defined(THREAD_SANITIZER)
55+
GTEST_SKIP() << "GPU tests are not supported In msan";
56+
#endif
57+
BenchmarkParams params = BenchmarkLiteRtModel::DefaultParams();
58+
params.Set<std::string>("graph", kModelPath);
59+
params.Set<std::string>("signature_to_run_for", kSignatureToRunFor);
60+
params.Set<bool>("use_cpu", false);
61+
params.Set<bool>("use_gpu", true);
62+
params.Set<bool>("require_full_delegation", true);
63+
64+
BenchmarkLiteRtModel benchmark = BenchmarkLiteRtModel(std::move(params));
65+
66+
EXPECT_EQ(benchmark.Run(), kTfLiteOk);
67+
}
68+
69+
TEST(BenchmarkLiteRtModelTest, GPUAccelerationWithProfiler) {
70+
// MSAN does not support GPU tests.
71+
#if defined(MEMORY_SANITIZER) || defined(THREAD_SANITIZER)
72+
GTEST_SKIP() << "GPU tests are not supported In msan";
73+
#endif
74+
BenchmarkParams params = BenchmarkLiteRtModel::DefaultParams();
75+
params.Set<std::string>("graph", kModelPath);
76+
params.Set<std::string>("signature_to_run_for", kSignatureToRunFor);
77+
params.Set<bool>("use_cpu", false);
78+
params.Set<bool>("use_gpu", true);
79+
params.Set<bool>("require_full_delegation", true);
80+
params.Set<bool>("use_profiler", true);
81+
82+
BenchmarkLiteRtModel benchmark = BenchmarkLiteRtModel(std::move(params));
83+
84+
EXPECT_EQ(benchmark.Run(), kTfLiteOk);
85+
}
86+
87+
} // namespace
88+
} // namespace benchmark
89+
} // namespace litert

‎litert/tools/benchmark_litert_model_test.cc‎

Lines changed: 0 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -70,41 +70,6 @@ TEST(BenchmarkLiteRtModelTest, GetModelSizeFromPathSucceeded) {
7070
EXPECT_GE(listener.results_.model_size_mb(), 0);
7171
}
7272

73-
TEST(BenchmarkLiteRtModelTest, GPUAcceleration) {
74-
// MSAN does not support GPU tests.
75-
#if defined(MEMORY_SANITIZER) || defined(THREAD_SANITIZER)
76-
GTEST_SKIP() << "GPU tests are not supported In msan";
77-
#endif
78-
BenchmarkParams params = BenchmarkLiteRtModel::DefaultParams();
79-
params.Set<std::string>("graph", kModelPath);
80-
params.Set<std::string>("signature_to_run_for", kSignatureToRunFor);
81-
params.Set<bool>("use_cpu", false);
82-
params.Set<bool>("use_gpu", true);
83-
params.Set<bool>("require_full_delegation", true);
84-
85-
BenchmarkLiteRtModel benchmark = BenchmarkLiteRtModel(std::move(params));
86-
87-
EXPECT_EQ(benchmark.Run(), kTfLiteOk);
88-
}
89-
90-
TEST(BenchmarkLiteRtModelTest, GPUAccelerationWithProfiler) {
91-
// MSAN does not support GPU tests.
92-
#if defined(MEMORY_SANITIZER) || defined(THREAD_SANITIZER)
93-
GTEST_SKIP() << "GPU tests are not supported In msan";
94-
#endif
95-
BenchmarkParams params = BenchmarkLiteRtModel::DefaultParams();
96-
params.Set<std::string>("graph", kModelPath);
97-
params.Set<std::string>("signature_to_run_for", kSignatureToRunFor);
98-
params.Set<bool>("use_cpu", false);
99-
params.Set<bool>("use_gpu", true);
100-
params.Set<bool>("require_full_delegation", true);
101-
params.Set<bool>("use_profiler", true);
102-
103-
BenchmarkLiteRtModel benchmark = BenchmarkLiteRtModel(std::move(params));
104-
105-
EXPECT_EQ(benchmark.Run(), kTfLiteOk);
106-
}
107-
10873
TEST(BenchmarkLiteRtModelTest, BenchmarkWithResultFilePath) {
10974
BenchmarkParams params = BenchmarkLiteRtModel::DefaultParams();
11075
params.Set<std::string>("graph", kModelPath);

0 commit comments

Comments
 (0)