mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
4edc8273eb
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15876 Build changes made it so some .so libraries are now registered after GlobalInit is called. Although this shouldn't be common, it also shouldn't be explicitly excluded. These changes allow for late Caffe2 registration, but also warn in that case. Reviewed By: kuttas Differential Revision: D13608186 fbshipit-source-id: 0ca7bcd32516d374077db0c2548cf8c28ccdd5f6
70 lines
1.8 KiB
C++
70 lines
1.8 KiB
C++
#include <iostream>
|
|
#include <memory>
|
|
|
|
#include <gtest/gtest.h>
|
|
#include "caffe2/core/init.h"
|
|
#include "caffe2/core/logging.h"
|
|
|
|
namespace caffe2 {
|
|
namespace {
|
|
bool gTestInitFunctionHasBeenRun = false;
|
|
bool gTestFailInitFunctionHasBeenRun = false;
|
|
|
|
bool TestInitFunction(int*, char***) {
|
|
gTestInitFunctionHasBeenRun = true;
|
|
return true;
|
|
}
|
|
|
|
bool TestFailInitFunction(int*, char***) {
|
|
gTestFailInitFunctionHasBeenRun = true;
|
|
return false;
|
|
}
|
|
|
|
REGISTER_CAFFE2_INIT_FUNCTION(
|
|
TestInitFunction,
|
|
&TestInitFunction,
|
|
"Just a test to see if GlobalInit invokes "
|
|
"registered functions correctly.");
|
|
|
|
int dummy_argc = 1;
|
|
const char* dummy_name = "foo";
|
|
char** dummy_argv = const_cast<char**>(&dummy_name);
|
|
} // namespace
|
|
|
|
TEST(InitTest, TestInitFunctionHasRun) {
|
|
caffe2::GlobalInit(&dummy_argc, &dummy_argv);
|
|
EXPECT_TRUE(gTestInitFunctionHasBeenRun);
|
|
EXPECT_FALSE(gTestFailInitFunctionHasBeenRun);
|
|
}
|
|
|
|
TEST(InitTest, CanRerunGlobalInit) {
|
|
caffe2::GlobalInit(&dummy_argc, &dummy_argv);
|
|
EXPECT_TRUE(caffe2::GlobalInit(&dummy_argc, &dummy_argv));
|
|
}
|
|
|
|
void LateRegisterInitFunction() {
|
|
::caffe2::InitRegisterer testInitFunc(
|
|
TestInitFunction, false, "This should succeed but warn");
|
|
}
|
|
|
|
void LateRegisterEarlyInitFunction() {
|
|
::caffe2::InitRegisterer testSecondInitFunc(
|
|
TestInitFunction, true, "This should fail for early init");
|
|
}
|
|
|
|
void LateRegisterFailInitFunction() {
|
|
::caffe2::InitRegisterer testSecondInitFunc(
|
|
TestFailInitFunction, false, "This should fail for failed init");
|
|
}
|
|
|
|
TEST(InitTest, FailLateRegisterInitFunction) {
|
|
caffe2::GlobalInit(&dummy_argc, &dummy_argv);
|
|
LateRegisterInitFunction();
|
|
EXPECT_THROW(LateRegisterEarlyInitFunction(), ::c10::Error);
|
|
EXPECT_THROW(LateRegisterFailInitFunction(), ::c10::Error);
|
|
EXPECT_TRUE(gTestInitFunctionHasBeenRun);
|
|
EXPECT_TRUE(gTestFailInitFunctionHasBeenRun);
|
|
}
|
|
|
|
} // namespace caffe2
|