| 158 | } |
| 159 | |
| 160 | void cpu_kernel( |
| 161 | const std::vector<Tensor>& inputs, const Param& params, |
| 162 | std::vector<Tensor>& outputs) { |
| 163 | ASSERT_TRUE(inputs.size() == 2); |
| 164 | ASSERT_TRUE(outputs.size() == 2); |
| 165 | ASSERT_TRUE(params["device"] == "x86"); |
| 166 | ASSERT_TRUE(params["scale_f"] == 2.12f); |
| 167 | ASSERT_TRUE(params["offset_i"] == 6); |
| 168 | ASSERT_TRUE(inputs[0].shape() == Shape({3, 4})); |
| 169 | ASSERT_TRUE(inputs[1].shape() == Shape({5, 6})); |
| 170 | ASSERT_TRUE(outputs[0].shape() == Shape({5, 6})); |
| 171 | ASSERT_TRUE(outputs[1].shape() == Shape({3, 4})); |
| 172 | ASSERT_TRUE(inputs[0].device() == "x86"); |
| 173 | ASSERT_TRUE(inputs[1].device() == "x86"); |
| 174 | ASSERT_TRUE(outputs[0].device() == "x86"); |
| 175 | ASSERT_TRUE(outputs[1].device() == "x86"); |
| 176 | |
| 177 | float scale_f = params["scale_f"].as<float>(); |
| 178 | int offset_i = params["offset_i"].as<int>(); |
| 179 | |
| 180 | for (size_t i = 0; i < 5 * 6; ++i) { |
| 181 | ASSERT_TRUE(inputs[1].data<float>()[i] == static_cast<float>(i)); |
| 182 | outputs[0].data<float>()[i] = inputs[1].data<float>()[i] * scale_f; |
| 183 | } |
| 184 | for (size_t i = 0; i < 3 * 4; ++i) { |
| 185 | ASSERT_TRUE(inputs[0].data<int>()[i] == static_cast<int>(i)); |
| 186 | outputs[1].data<int>()[i] = inputs[0].data<int>()[i] + offset_i; |
| 187 | } |
| 188 | } |
| 189 | |
| 190 | TEST(TestCustomOp, TestCustomOpCompute) { |
| 191 | std::shared_ptr<CustomOp> op = |