@@ -49,7 +49,7 @@ Kernel* registered_kernels = reinterpret_cast<Kernel*>(registered_kernels_data);
49
49
size_t num_registered_kernels = 0 ;
50
50
51
51
// Registers the kernels, but may return an error.
52
- Error register_kernels_internal (const Span<const Kernel> kernels) {
52
+ Error register_kernels_internal (const Span<const Kernel> kernels, ErrorHandler errorHandler ) {
53
53
// Operator registration happens in static initialization time before or after
54
54
// PAL init, so call it here. It is safe to call multiple times.
55
55
::et_pal_init ();
@@ -74,12 +74,19 @@ Error register_kernels_internal(const Span<const Kernel> kernels) {
74
74
ET_LOG (Error, " %s" , kernels[i].name_ );
75
75
ET_LOG_KERNEL_KEY (kernels[i].kernel_key_ );
76
76
}
77
+
78
+ if (errorHandler != nullptr ) {
79
+ return errorHandler (Error::RegistrationExceedingMaxKernels);
80
+ }
81
+
77
82
return Error::RegistrationExceedingMaxKernels;
78
83
}
79
84
// for debugging purpose
80
85
ET_UNUSED const char * lib_name =
81
86
et_pal_get_shared_library_name (kernels.data ());
82
87
88
+ Error err = Error::Ok;
89
+
83
90
for (const auto & kernel : kernels) {
84
91
// Linear search. This is fine if the number of kernels is small.
85
92
for (size_t i = 0 ; i < num_registered_kernels; i++) {
@@ -88,24 +95,33 @@ Error register_kernels_internal(const Span<const Kernel> kernels) {
88
95
kernel.kernel_key_ == k.kernel_key_ ) {
89
96
ET_LOG (Error, " Re-registering %s, from %s" , k.name_ , lib_name);
90
97
ET_LOG_KERNEL_KEY (k.kernel_key_ );
91
- return Error::RegistrationAlreadyRegistered;
98
+ err = Error::RegistrationAlreadyRegistered;
99
+ continue ;
92
100
}
93
101
}
102
+
94
103
registered_kernels[num_registered_kernels++] = kernel;
95
104
}
96
- ET_LOG (
97
- Debug,
98
- " Successfully registered all kernels from shared library: %s" ,
99
- lib_name);
100
105
101
- return Error::Ok;
106
+ if (errorHandler != nullptr ) {
107
+ err = errorHandler (err);
108
+ }
109
+
110
+ if (err == Error::Ok) {
111
+ ET_LOG (
112
+ Debug,
113
+ " Successfully registered all kernels from shared library: %s" ,
114
+ lib_name);
115
+ }
116
+
117
+ return err;
102
118
}
103
119
104
120
} // namespace
105
121
106
122
// Registers the kernels, but panics if an error occurs. Always returns Ok.
107
- Error register_kernels (const Span<const Kernel> kernels) {
108
- Error success = register_kernels_internal (kernels);
123
+ Error register_kernels (const Span<const Kernel> kernels, ErrorHandler errorHandler ) {
124
+ Error success = register_kernels_internal (kernels, errorHandler );
109
125
if (success == Error::RegistrationAlreadyRegistered ||
110
126
success == Error::RegistrationExceedingMaxKernels) {
111
127
ET_CHECK_MSG (
0 commit comments