Skip to content

Commit fda5ee9

Browse files
terryheocopybara-github
terryheo
authored and
copybara-github
committed
Add Jet Accelerator
This is an experimental legacy TFLite Jet based GPU Accelerator. It doesn't support GPU TensorBuffer binding like other GPU Accelerator but uses CPU memory TensorBuffer. Also added benchmark_jet, gpu_numerics_check_jet LiteRT-PiperOrigin-RevId: 820808902
1 parent d20b613 commit fda5ee9

File tree

3 files changed

+217
-0
lines changed

3 files changed

+217
-0
lines changed

‎litert/cc/BUILD‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -602,6 +602,63 @@ litert_device_test(
602602
# ],
603603
# )
604604
#
605+
# # Should use `--copt=-DCL_DELEGATE_NO_GL` to build for Desktop.
606+
# cc_test(
607+
# name = "litert_compiled_model_jetgpu_test",
608+
# srcs = [":litert_compiled_model_jetgpu_test.cc"],
609+
# data = [
610+
# "//litert/test:mlir_test_data",
611+
# "//litert/test:tflite_test_data",
612+
# ],
613+
# tags = [
614+
# "no_oss",
615+
# "notap",
616+
# "requires-gpu-nvidia",
617+
# ],
618+
# deps = [
619+
# "@com_google_googletest//:gtest_main",
620+
# "@com_google_absl//absl/container:flat_hash_map",
621+
# "@com_google_absl//absl/debugging:leak_check",
622+
# "@com_google_absl//absl/log:absl_log",
623+
# "@com_google_absl//absl/strings:string_view",
624+
# "@com_google_absl//absl/types:span",
625+
# "//litert/c:litert_common",
626+
# "//litert/c:litert_environment_options",
627+
# "//litert/c:litert_event",
628+
# "//litert/c:litert_event_type",
629+
# "//litert/c:litert_profiler_event",
630+
# "//litert/c:litert_tensor_buffer",
631+
# "//litert/c:litert_tensor_buffer_types",
632+
# "//litert/cc:litert_compiled_model",
633+
# "//litert/cc:litert_element_type",
634+
# "//litert/cc:litert_environment",
635+
# "//litert/cc:litert_event",
636+
# "//litert/cc:litert_expected",
637+
# "//litert/cc:litert_layout",
638+
# "//litert/cc:litert_macros",
639+
# "//litert/cc:litert_model",
640+
# "//litert/cc:litert_options",
641+
# "//litert/cc:litert_profiler",
642+
# "//litert/cc:litert_tensor_buffer",
643+
# "//litert/cc/internal:litert_platform_support",
644+
# "//litert/cc/options:litert_gpu_options",
645+
# "//litert/cc/options:litert_runtime_options",
646+
# "//litert/runtime/accelerators/gpu/google:jet_gpu_accelerator", # buildcleaner: keep
647+
# "//litert/test:common",
648+
# "//litert/test:matchers",
649+
# "//litert/test:simple_model",
650+
# "@opencl_headers",
651+
# ] + select({
652+
# "@org_tensorflow//tensorflow:android": [
653+
# "//tflite/delegates/gpu/cl:cl_device",
654+
# "//tflite/delegates/gpu/cl:gl_interop",
655+
# "//tflite/delegates/gpu/cl:opencl_wrapper",
656+
# "//tflite/delegates/gpu/gl:egl_environment",
657+
# ],
658+
# "//conditions:default": [],
659+
# }),
660+
# )
661+
#
605662
# copybara:uncomment_end
606663

607664
cc_library(

‎litert/cc/litert_compiled_model_jetgpu_test.cc‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
// Copyright 2025 Google LLC.
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+
#include <cstring>
16+
#include <tuple>
17+
#include <utility>
18+
19+
#include <gmock/gmock.h>
20+
#include <gtest/gtest.h>
21+
#include "absl/debugging/leak_check.h" // from @com_google_absl
22+
#include "absl/log/absl_log.h" // from @com_google_absl
23+
#include "absl/strings/string_view.h" // from @com_google_absl
24+
#include "absl/types/span.h" // from @com_google_absl
25+
#include "litert/c/litert_common.h"
26+
#include "litert/c/litert_tensor_buffer_types.h"
27+
#include "litert/cc/litert_compiled_model.h"
28+
#include "litert/cc/litert_environment.h"
29+
#include "litert/cc/litert_expected.h"
30+
#include "litert/cc/litert_macros.h"
31+
#include "litert/cc/litert_model.h"
32+
#include "litert/cc/litert_options.h"
33+
#include "litert/cc/litert_tensor_buffer.h"
34+
#include "litert/cc/options/litert_gpu_options.h"
35+
#include "litert/test/common.h"
36+
#include "litert/test/matchers.h"
37+
#include "litert/test/testdata/simple_model_test_vectors.h"
38+
39+
using testing::FloatNear;
40+
using testing::Pointwise;
41+
42+
namespace litert {
43+
namespace {
44+
45+
using TestParams = std::tuple<LiteRtDelegatePrecision>;
46+
47+
Expected<Options> CreateGpuOptions(const TestParams& params) {
48+
LITERT_ASSIGN_OR_RETURN(auto gpu_options, GpuOptions::Create());
49+
LITERT_RETURN_IF_ERROR(gpu_options.SetDelegatePrecision(std::get<0>(params)));
50+
51+
LITERT_ASSIGN_OR_RETURN(litert::Options options, Options::Create());
52+
options.SetHardwareAccelerators(kLiteRtHwAcceleratorGpu);
53+
options.AddOpaqueOptions(std::move(gpu_options));
54+
return std::move(options);
55+
}
56+
57+
class ParameterizedTest : public ::testing::TestWithParam<TestParams> {};
58+
59+
TEST_P(ParameterizedTest, Basic) {
60+
// To workaround the memory leak in Nvidia's driver
61+
absl::LeakCheckDisabler disable_leak_check;
62+
63+
LITERT_ASSERT_OK_AND_ASSIGN(
64+
auto model,
65+
Model::CreateFromFile(testing::GetTestFilePath(kModelFileName)));
66+
67+
auto env = litert::Environment::Create({});
68+
ASSERT_TRUE(env);
69+
70+
LITERT_ASSERT_OK_AND_ASSIGN(auto options, CreateGpuOptions(GetParam()));
71+
LITERT_ASSERT_OK_AND_ASSIGN(auto compiled_model,
72+
CompiledModel::Create(*env, model, options));
73+
74+
LITERT_ASSERT_OK_AND_ASSIGN(auto signatures, model.GetSignatures());
75+
EXPECT_EQ(signatures.size(), 1);
76+
77+
auto signature_key = signatures[0].Key();
78+
EXPECT_EQ(signature_key, Model::DefaultSignatureKey());
79+
size_t signature_index = 0;
80+
81+
LITERT_ASSERT_OK_AND_ASSIGN(
82+
auto input_buffers, compiled_model.CreateInputBuffers(signature_index));
83+
84+
LITERT_ASSERT_OK_AND_ASSIGN(
85+
auto output_buffers, compiled_model.CreateOutputBuffers(signature_index));
86+
87+
// Fill model inputs.
88+
auto input_names = signatures[0].InputNames();
89+
EXPECT_EQ(input_names.size(), 2);
90+
EXPECT_EQ(input_names.at(0), "arg0");
91+
EXPECT_EQ(input_names.at(1), "arg1");
92+
LITERT_ASSERT_OK_AND_ASSIGN(auto buffer0_type, input_buffers[0].BufferType());
93+
EXPECT_EQ(buffer0_type, kLiteRtTensorBufferTypeHostMemory);
94+
ASSERT_TRUE(input_buffers[0].Write<float>(
95+
absl::MakeConstSpan(kTestInput0Tensor, kTestInput0Size)));
96+
LITERT_ASSERT_OK_AND_ASSIGN(auto buffer1_type, input_buffers[0].BufferType());
97+
EXPECT_EQ(buffer1_type, kLiteRtTensorBufferTypeHostMemory);
98+
ASSERT_TRUE(input_buffers[1].Write<float>(
99+
absl::MakeConstSpan(kTestInput1Tensor, kTestInput1Size)));
100+
101+
// Execute model.
102+
compiled_model.Run(signature_index, input_buffers, output_buffers);
103+
104+
// Check model output.
105+
auto output_names = signatures[0].OutputNames();
106+
EXPECT_EQ(output_names.size(), 1);
107+
EXPECT_EQ(output_names.at(0), "tfl.add");
108+
LITERT_ASSERT_OK_AND_ASSIGN(auto output_type, input_buffers[0].BufferType());
109+
EXPECT_EQ(output_type, kLiteRtTensorBufferTypeHostMemory);
110+
{
111+
auto lock_and_addr = litert::TensorBufferScopedLock::Create<const float>(
112+
output_buffers[0], TensorBuffer::LockMode::kRead);
113+
ASSERT_TRUE(lock_and_addr);
114+
auto output = absl::MakeSpan(lock_and_addr->second, kTestOutputSize);
115+
for (auto i = 0; i < kTestOutputSize; ++i) {
116+
ABSL_LOG(INFO) << "Result: " << output[i] << "\t" << kTestOutputTensor[i];
117+
}
118+
EXPECT_THAT(output, Pointwise(FloatNear(1e-5), kTestOutputTensor));
119+
}
120+
}
121+
122+
INSTANTIATE_TEST_SUITE_P(
123+
CompiledModelWebGpuTest, ParameterizedTest,
124+
::testing::Combine(::testing::ValuesIn<LiteRtDelegatePrecision>({
125+
kLiteRtDelegatePrecisionDefault,
126+
kLiteRtDelegatePrecisionFp16,
127+
kLiteRtDelegatePrecisionFp32,
128+
})));
129+
130+
} // namespace
131+
} // namespace litert

‎litert/tools/BUILD‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -534,6 +534,21 @@ cc_binary(
534534
],
535535
)
536536

537+
cc_binary(
538+
name = "gpu_numerics_check_jet",
539+
srcs = ["gpu_numerics_check.cc"],
540+
tags = [
541+
"notap",
542+
"requires-gpu-nvidia",
543+
],
544+
deps = NUMERICS_CHECK_DEPS + [
545+
"//litert/cc/options:litert_gpu_options",
546+
# copybara:uncomment_begin(google-only)
547+
# "//litert/runtime/accelerators/gpu/google:jet_gpu_accelerator", # buildcleaner: keep
548+
# copybara:uncomment_end
549+
],
550+
)
551+
537552
cc_binary(
538553
name = "analyze_model_main",
539554
srcs = ["analyze_model_main.cc"],
@@ -593,6 +608,20 @@ cc_binary(
593608
],
594609
)
595610

611+
cc_binary(
612+
name = "benchmark_model_jet",
613+
tags = [
614+
"notap",
615+
"requires-gpu-nvidia",
616+
],
617+
deps = [
618+
":benchmark_model_main",
619+
# copybara:uncomment_begin(google-only)
620+
# "//litert/runtime/accelerators/gpu/google:jet_gpu_accelerator", # buildcleaner: keep
621+
# copybara:uncomment_end
622+
],
623+
)
624+
596625
cc_library(
597626
name = "benchmark_stripped_litert_model",
598627
srcs = ["benchmark_stripped_litert_model.cc"],

0 commit comments

Comments
 (0)