| 223 | |
| 224 | template <typename FunctionType, typename ExecType> |
| 225 | void CheckAddDispatch(FunctionType* func, ExecType exec) { |
| 226 | using KernelType = typename FunctionType::KernelType; |
| 227 | |
| 228 | ASSERT_EQ(0, func->num_kernels()); |
| 229 | ASSERT_EQ(0, func->kernels().size()); |
| 230 | |
| 231 | std::vector<InputType> in_types1 = {int32(), int32()}; |
| 232 | OutputType out_type1 = int32(); |
| 233 | |
| 234 | ASSERT_OK(func->AddKernel(in_types1, out_type1, exec)); |
| 235 | ASSERT_OK(func->AddKernel({int32(), int8()}, int32(), exec)); |
| 236 | |
| 237 | // Duplicate sig is okay |
| 238 | ASSERT_OK(func->AddKernel(in_types1, out_type1, exec)); |
| 239 | |
| 240 | // Add a kernel |
| 241 | KernelType kernel({float64(), float64()}, float64(), exec); |
| 242 | ASSERT_OK(func->AddKernel(kernel)); |
| 243 | |
| 244 | ASSERT_EQ(4, func->num_kernels()); |
| 245 | ASSERT_EQ(4, func->kernels().size()); |
| 246 | |
| 247 | // Try adding some invalid kernels |
| 248 | ASSERT_RAISES(Invalid, func->AddKernel({}, int32(), exec)); |
| 249 | ASSERT_RAISES(Invalid, func->AddKernel({int32()}, int32(), exec)); |
| 250 | ASSERT_RAISES(Invalid, func->AddKernel({int8(), int8(), int8()}, int32(), exec)); |
| 251 | |
| 252 | // Add valid and invalid kernel using kernel struct directly |
| 253 | KernelType valid_kernel({boolean(), boolean()}, boolean(), exec); |
| 254 | ASSERT_OK(func->AddKernel(valid_kernel)); |
| 255 | |
| 256 | KernelType invalid_kernel({boolean()}, boolean(), exec); |
| 257 | ASSERT_RAISES(Invalid, func->AddKernel(invalid_kernel)); |
| 258 | |
| 259 | ASSERT_OK_AND_ASSIGN(const Kernel* dispatched, func->DispatchExact({int32(), int32()})); |
| 260 | KernelSignature expected_sig(in_types1, out_type1); |
| 261 | ASSERT_TRUE(dispatched->signature->Equals(expected_sig)); |
| 262 | |
| 263 | // No kernel available |
| 264 | ASSERT_RAISES(NotImplemented, func->DispatchExact({utf8(), utf8()})); |
| 265 | |
| 266 | // Wrong arity |
| 267 | ASSERT_RAISES(Invalid, func->DispatchExact({})); |
| 268 | ASSERT_RAISES(Invalid, func->DispatchExact({int32(), int32(), int32()})); |
| 269 | } |
| 270 | |
| 271 | TEST(ScalarVectorFunction, DispatchExact) { |
| 272 | ScalarFunction func1("scalar_test", Arity::Binary(), /*doc=*/FunctionDoc::Empty()); |
no test coverage detected