Skip to content

Commit a2f59e9

Browse files
shoumikhinDannyYuyang-quic
authored andcommitted
MethodMeta API to query of backends number and names (pytorch#9440)
1 parent 233670f commit a2f59e9

File tree

3 files changed

+58
-0
lines changed

3 files changed

+58
-0
lines changed

runtime/executor/method_meta.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,19 @@ bool MethodMeta::uses_backend(const char* backend_name) const {
239239
return false;
240240
}
241241

242+
size_t MethodMeta::num_backends() const {
243+
const auto delegates = s_plan_->delegates();
244+
return delegates ? delegates->size() : 0;
245+
}
246+
247+
const char* MethodMeta::get_backend_name(size_t index) const {
248+
const auto delegates = s_plan_->delegates();
249+
if (delegates && index < delegates->size()) {
250+
return delegates->Get(index)->id()->c_str();
251+
}
252+
return nullptr;
253+
}
254+
242255
size_t MethodMeta::num_instructions() const {
243256
const auto chains = s_plan_->chains();
244257
if (chains == nullptr) {

runtime/executor/method_meta.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,22 @@ class MethodMeta final {
193193
*/
194194
bool uses_backend(const char* backend_name) const;
195195

196+
/**
197+
* Get the number of backends used in this method.
198+
*
199+
* @returns The total number of backend names.
200+
*/
201+
size_t num_backends() const;
202+
203+
/**
204+
* Get the backend name at the given index.
205+
*
206+
* @param[in] index The index of the backend name.
207+
* @returns The backend name as a C-style string, or nullptr if the index is
208+
* invalid.
209+
*/
210+
const char* get_backend_name(size_t index) const;
211+
196212
/**
197213
* Get the number of instructions in this method.
198214
*

runtime/executor/test/backend_integration_test.cpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,35 @@ TEST_P(BackendIntegrationTest, BasicInitSucceeds) {
347347
EXPECT_EQ(method_res.error(), Error::Ok);
348348
}
349349

350+
TEST_P(BackendIntegrationTest, GetBackendNamesSuccess) {
351+
// Load the program from file.
352+
Result<FileDataLoader> loader = FileDataLoader::from(program_path());
353+
ASSERT_EQ(loader.error(), Error::Ok);
354+
355+
Result<Program> program = Program::load(&loader.get());
356+
ASSERT_EQ(program.error(), Error::Ok);
357+
358+
// Get method metadata for the "forward" method.
359+
auto method_meta = program->method_meta("forward");
360+
361+
// Ensure the StubBackend is used.
362+
EXPECT_TRUE(method_meta->uses_backend(StubBackend::kName));
363+
364+
// Retrieve the number of backends.
365+
size_t num_backends = method_meta->num_backends();
366+
EXPECT_GT(num_backends, 0u);
367+
368+
// Iterate through each backend and verify its name.
369+
for (size_t i = 0; i < num_backends; ++i) {
370+
const char* name = method_meta->get_backend_name(i);
371+
EXPECT_NE(name, nullptr);
372+
// For this test, we expect that the only backend is StubBackend.
373+
EXPECT_STREQ(name, StubBackend::kName);
374+
}
375+
// Check that an out-of-range index returns nullptr.
376+
EXPECT_EQ(method_meta->get_backend_name(num_backends), nullptr);
377+
}
378+
350379
TEST_P(BackendIntegrationTest, FreeingProcessedBufferSucceeds) {
351380
// Install an init() implementation that frees its processed buffer, and lets
352381
// us know that it was actually called by setting init_called.

0 commit comments

Comments
 (0)