intel-extension-for-pytorch

1
#include <torch/torch.h>
2
#include "csrc/cpu/runtime/CPUPool.h"
3
#include "csrc/cpu/runtime/Task.h"
4
#include "csrc/cpu/runtime/TaskExecutor.h"
5
#include "gtest/gtest.h"
6

7
#define ASSERT_VARIABLE_EQ(a, b) ASSERT_TRUE(torch::allclose((a), (b)))
8
#define EXPECT_VARIABLE_EQ(a, b) EXPECT_TRUE(torch::allclose((a), (b)))
9

10
TEST(TestRuntimeAPI, TestMainThreadCoreBind) {
11
  // 1. Get the default thread affinity information in main thread.
12
  // 2. Set the new thread affinity information in main thread.
13
  // 3. Run the function in main thread.
14
  // 4. Restore the default thread affinity information in main thread.
15
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
16
    GTEST_SKIP()
17
        << "Skip TestRuntimeAPI::TestMainThreadCoreBind. Didn't preload IOMP.";
18
  }
19
  at::Tensor input_tensor = at::rand({100, 8276});
20
  // Get the reference result.
21
  auto res_ref = at::softmax(input_tensor, -1);
22
  // Get current cpu_pool information.
23
  torch_ipex::runtime::CPUPool previous_cpu_pool =
24
      torch_ipex::runtime::get_cpu_pool_from_mask_affinity();
25
  // Ping CPU Cores.
26
  std::vector<int32_t> cpu_core_list({0});
27
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
28
  torch_ipex::runtime::_pin_cpu_cores(cpu_pool);
29
  auto res = at::softmax(input_tensor, -1);
30
  ASSERT_VARIABLE_EQ(res, res_ref);
31
  // restore the cpu pool information.
32
  torch_ipex::runtime::set_mask_affinity_from_cpu_pool(previous_cpu_pool);
33
}
34

35
TEST(TestRuntimeAPI, TestMainThreadCoreBindWithCPUPool) {
36
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
37
    GTEST_SKIP()
38
        << "Skip TestRuntimeAPI::TestMainThreadCoreBindWithCPUPool. Didn't preload IOMP.";
39
  }
40
  at::Tensor input_tensor = at::rand({100, 8276});
41
  std::vector<int32_t> cpu_core_list({0});
42
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
43
  {
44
    torch_ipex::runtime::WithCPUPool with_cpu_pool(std::move(cpu_pool));
45
    auto res = at::softmax(input_tensor, -1);
46
  }
47
  auto res_ = at::softmax(input_tensor, -1);
48
}
49

50
TEST(TestRuntimeTaskAPI, TestTaskAPINativeTorchOperation) {
51
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
52
    GTEST_SKIP()
53
        << "Skip TestRuntimeTaskAPI::TestTaskAPINativeTorchOperation. Didn't preload IOMP.";
54
  }
55
  std::vector<int32_t> cpu_core_list({0});
56
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
57
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
58
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
59
  at::Tensor input_tensor = at::rand({100, 8276});
60
  // Get the reference result
61
  auto res_ref = at::softmax(input_tensor, -1);
62
  // Create the task
63
  torch_ipex::runtime::Task<
64
      at::Tensor (*)(const at::Tensor&, int64_t, c10::optional<at::ScalarType>),
65
      const at::Tensor&,
66
      int64_t,
67
      c10::optional<at::ScalarType>&>
68
      task(at::softmax, task_executor);
69
  c10::optional<at::ScalarType> dtype = c10::nullopt;
70
  // or
71
  // c10::optional<at::ScalarType> dtype = input_tensor.scalar_type();
72
  auto res_future = task(input_tensor, -1, dtype);
73
  auto res = res_future.get();
74

75
  // Test Rvalue Input Tensor
76
  auto res_future_rinput = task(std::move(input_tensor), -1, dtype);
77
  auto res_rinput = res_future_rinput.get();
78

79
  // Test Const Input Tensor
80
  const at::Tensor input_tensor2 = at::rand({100, 8276});
81
  auto res_ref_const_input = at::softmax(input_tensor2, -1);
82
  auto res_future_const_input = task(input_tensor2, -1, dtype);
83
  auto res_const_input = res_future_const_input.get();
84

85
  // Assert the result
86
  ASSERT_VARIABLE_EQ(res, res_ref);
87
  ASSERT_VARIABLE_EQ(res_rinput, res_ref);
88
  ASSERT_VARIABLE_EQ(res_const_input, res_ref_const_input);
89
}
90

91
TEST(TestRuntimeTaskAPI, TestTaskAPILambdaFunction) {
92
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
93
    GTEST_SKIP()
94
        << "Skip TestRuntimeTaskAPI::TestTaskAPILambdaFunction. Didn't preload IOMP.";
95
  }
96
  std::vector<int32_t> cpu_core_list({0});
97
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
98
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
99
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
100
  at::Tensor input_tensor = at::rand({100, 8276});
101
  // Get the reference result
102
  auto res_ref = at::softmax(input_tensor, -1);
103
  // Create the task
104
  torch_ipex::runtime::
105
      Task<at::Tensor (*)(const at::Tensor&), const at::Tensor&>
106
          task(
107
              [](const at::Tensor& input) -> at::Tensor {
108
                return at::softmax(input, -1);
109
              },
110
              task_executor);
111
  auto res_future = task(input_tensor);
112
  auto res = res_future.get();
113
  // Assert the result
114
  ASSERT_VARIABLE_EQ(res, res_ref);
115
}
116

117
at::Tensor taskfunction_native_input(at::Tensor input) {
118
  at::Tensor output;
119
  output = at::softmax(input, -1);
120
  return output;
121
}
122

123
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionNativeInput) {
124
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
125
    GTEST_SKIP()
126
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionNativeInput. Didn't preload IOMP.";
127
  }
128
  std::vector<int32_t> cpu_core_list({0});
129
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
130
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
131
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
132
  at::Tensor input_tensor = at::rand({100, 8276});
133
  // Get the reference result
134
  auto res_ref = taskfunction_native_input(input_tensor);
135
  // Create the task
136
  torch_ipex::runtime::Task<at::Tensor (*)(at::Tensor), at::Tensor> task(
137
      taskfunction_native_input, task_executor);
138
  auto res_future = task(std::move(input_tensor));
139
  auto res = res_future.get();
140
  // Assert the result
141
  ASSERT_VARIABLE_EQ(res, res_ref);
142
}
143

144
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionNativeInputLValue) {
145
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
146
    GTEST_SKIP()
147
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionNativeInput. Didn't preload IOMP.";
148
  }
149
  std::vector<int32_t> cpu_core_list({0});
150
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
151
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
152
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
153
  at::Tensor input_tensor = at::rand({100, 8276});
154
  // Get the reference result
155
  auto res_ref = taskfunction_native_input(input_tensor);
156
  // Create the task
157
  torch_ipex::runtime::Task<at::Tensor (*)(at::Tensor), at::Tensor&> task(
158
      taskfunction_native_input, task_executor);
159
  auto res_future = task(input_tensor);
160
  auto res = res_future.get();
161
  // Assert the result
162
  ASSERT_VARIABLE_EQ(res, res_ref);
163
}
164

165
at::Tensor taskfunction_lvalue_reference(at::Tensor& input) {
166
  at::Tensor output;
167
  output = at::softmax(input, -1);
168
  return output;
169
}
170

171
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionLValueReference) {
172
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
173
    GTEST_SKIP()
174
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionLValueReference. Didn't preload IOMP.";
175
  }
176
  std::vector<int32_t> cpu_core_list({0});
177
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
178
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
179
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
180
  at::Tensor input_tensor = at::rand({100, 8276});
181
  // Get the reference result
182
  auto res_ref = taskfunction_lvalue_reference(input_tensor);
183
  // Create the task
184
  torch_ipex::runtime::Task<at::Tensor (*)(at::Tensor&), at::Tensor&> task(
185
      taskfunction_lvalue_reference, task_executor);
186
  auto res_future = task(input_tensor);
187
  auto res = res_future.get();
188
  // Assert the result
189
  ASSERT_VARIABLE_EQ(res, res_ref);
190
}
191

192
at::Tensor taskfunction_const_lvalue_reference(const at::Tensor& input) {
193
  at::Tensor output;
194
  output = at::softmax(input, -1);
195
  return output;
196
}
197

198
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionConstLValueReference) {
199
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
200
    GTEST_SKIP()
201
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionConstLValueReference. Didn't preload IOMP.";
202
  }
203
  std::vector<int32_t> cpu_core_list({0});
204
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
205
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
206
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
207
  at::Tensor input_tensor = at::rand({100, 8276});
208
  // Get the reference result
209
  auto res_ref = taskfunction_const_lvalue_reference(input_tensor);
210
  // Create the task
211
  torch_ipex::runtime::
212
      Task<at::Tensor (*)(const at::Tensor&), const at::Tensor&>
213
          task(taskfunction_const_lvalue_reference, task_executor);
214
  auto res_future = task(input_tensor);
215
  auto res = res_future.get();
216
  // Assert the result
217
  ASSERT_VARIABLE_EQ(res, res_ref);
218
}
219

220
at::Tensor taskfunction_rvalue_reference(at::Tensor&& input) {
221
  at::Tensor output;
222
  output = at::softmax(input, -1);
223
  return output;
224
}
225

226
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionRvalueReference) {
227
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
228
    GTEST_SKIP()
229
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionRvalueReference. Didn't preload IOMP.";
230
  }
231
  std::vector<int32_t> cpu_core_list({0});
232
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
233
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
234
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
235
  at::Tensor input_tensor = at::rand({100, 8276});
236
  at::Tensor input_tensor2 = input_tensor;
237
  // Get the reference result
238
  auto res_ref = taskfunction_rvalue_reference(std::move(input_tensor));
239
  // Create the task
240
  torch_ipex::runtime::Task<at::Tensor (*)(at::Tensor &&), at::Tensor&&> task(
241
      taskfunction_rvalue_reference, task_executor);
242
  auto res_future = task(std::move(input_tensor2));
243
  auto res = res_future.get();
244
  // Assert the result
245
  ASSERT_VARIABLE_EQ(res, res_ref);
246
}
247

248
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>
249
taskfunction_mix_lvalue_rvalue_reference(
250
    at::Tensor input1,
251
    at::Tensor& input2,
252
    const at::Tensor& input3,
253
    at::Tensor&& input4) {
254
  at::Tensor output1 = at::softmax(input1, -1);
255
  at::Tensor output2 = at::softmax(input2, -1);
256
  at::Tensor output3 = at::softmax(input3, -1);
257
  at::Tensor output4 = at::softmax(input4, -1);
258
  return std::make_tuple(output1, output2, output3, output4);
259
}
260

261
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionMixLvalueRvalueReference) {
262
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
263
    GTEST_SKIP()
264
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionMixLvalueRvalueReference. Didn't preload IOMP.";
265
  }
266
  std::vector<int32_t> cpu_core_list({0});
267
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
268
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
269
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
270
  at::Tensor input_tensor = at::rand({100, 8276});
271
  at::Tensor input_tensor2 = input_tensor;
272
  at::Tensor input_tensor3 = input_tensor;
273
  at::Tensor input_tensor4 = input_tensor;
274
  // Get the reference result
275
  auto res_ref = taskfunction_mix_lvalue_rvalue_reference(
276
      std::move(input_tensor),
277
      input_tensor2,
278
      input_tensor2,
279
      std::move(input_tensor2));
280
  // Create the task
281
  torch_ipex::runtime::Task<
282
      std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> (*)(
283
          at::Tensor, at::Tensor&, const at::Tensor&, at::Tensor&&),
284
      at::Tensor,
285
      at::Tensor&,
286
      const at::Tensor&,
287
      at::Tensor&&>
288
      task(taskfunction_mix_lvalue_rvalue_reference, task_executor);
289
  auto res_future = task(
290
      std::move(input_tensor3),
291
      input_tensor4,
292
      input_tensor4,
293
      std::move(input_tensor4));
294
  auto res = res_future.get();
295
  // Assert the result
296
  ASSERT_VARIABLE_EQ(std::get<0>(res), std::get<0>(res_ref));
297
  ASSERT_VARIABLE_EQ(std::get<1>(res), std::get<1>(res_ref));
298
  ASSERT_VARIABLE_EQ(std::get<2>(res), std::get<2>(res_ref));
299
  ASSERT_VARIABLE_EQ(std::get<3>(res), std::get<3>(res_ref));
300
}
301

302
at::Tensor taskfunction_input_vector(std::vector<at::Tensor>& inputs) {
303
  at::Tensor output;
304
  output = at::softmax(inputs[0], -1);
305
  return output;
306
}
307

308
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionInputVectorTensor) {
309
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
310
    GTEST_SKIP()
311
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionInputVectorTensor. Didn't preload IOMP.";
312
  }
313
  std::vector<int32_t> cpu_core_list({0});
314
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
315
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
316
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
317
  at::Tensor input_tensor = at::rand({100, 8276});
318
  std::vector<at::Tensor> input_tenosrs;
319
  input_tenosrs.emplace_back(input_tensor);
320
  // Get the reference result
321
  auto res_ref = taskfunction_input_vector(input_tenosrs);
322
  // Create the task
323
  torch_ipex::runtime::
324
      Task<at::Tensor (*)(std::vector<at::Tensor>&), std::vector<at::Tensor>&>
325
          task(taskfunction_input_vector, task_executor);
326
  auto res_future = task(input_tenosrs);
327
  auto res = res_future.get();
328
  // Assert the result
329
  ASSERT_VARIABLE_EQ(res, res_ref);
330
}
331

332
at::Tensor& taskfunction_input_reference_output_lvalue_reference(
333
    at::Tensor& input,
334
    at::Tensor& output) {
335
  output = at::softmax(input, -1);
336
  return output;
337
}
338

339
TEST(TestRuntimeTaskAPI, TestTaskAPICPPFunctionOutputTensorLValueReference) {
340
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
341
    GTEST_SKIP()
342
        << "Skip TestRuntimeTaskAPI::TestTaskAPICPPFunctionOutputTensorLValueReference. Didn't preload IOMP.";
343
  }
344
  std::vector<int32_t> cpu_core_list({0});
345
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
346
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
347
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
348
  at::Tensor input_tensor = at::rand({100, 8276});
349
  at::Tensor output_tensor;
350
  at::Tensor output_tensor2;
351
  // Get the reference result
352
  auto res_ref = taskfunction_input_reference_output_lvalue_reference(
353
      input_tensor, output_tensor);
354
  // Create the task
355
  torch_ipex::runtime::
356
      Task<at::Tensor& (*)(at::Tensor&, at::Tensor&), at::Tensor&, at::Tensor&>
357
          task(
358
              taskfunction_input_reference_output_lvalue_reference,
359
              task_executor);
360
  auto res_future = task(input_tensor, output_tensor2);
361
  auto res = res_future.get();
362
  // Assert the result
363
  ASSERT_VARIABLE_EQ(res, res_ref);
364
  ASSERT_VARIABLE_EQ(output_tensor, res_ref);
365
  ASSERT_VARIABLE_EQ(output_tensor2, res_ref);
366
  ASSERT_VARIABLE_EQ(output_tensor2, output_tensor);
367
}
368

369
TEST(TestRuntimeTaskAPI, TestTaskAPIMultiTasksSameTensorInput) {
370
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
371
    GTEST_SKIP()
372
        << "Skip TestRuntimeTaskAPI::TestTaskAPIMultiTasksSameTensorInput. Didn't preload IOMP.";
373
  }
374
  std::vector<int32_t> cpu_core_list({0});
375
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
376
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
377
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
378

379
  std::vector<int32_t> cpu_core_list2({1});
380
  torch_ipex::runtime::CPUPool cpu_pool2(cpu_core_list2);
381
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor2 =
382
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool2);
383

384
  at::Tensor input_tensor = at::rand({100, 8276});
385
  // Get the reference result
386
  auto res_ref = at::softmax(input_tensor, -1);
387
  // Create the task
388
  torch_ipex::runtime::
389
      Task<at::Tensor (*)(const at::Tensor&), const at::Tensor&>
390
          task(taskfunction_const_lvalue_reference, task_executor);
391

392
  torch_ipex::runtime::
393
      Task<at::Tensor (*)(const at::Tensor&), const at::Tensor&>
394
          task2(taskfunction_const_lvalue_reference, task_executor2);
395

396
  auto res_future = task(input_tensor);
397
  auto res_future2 = task2(input_tensor);
398
  auto res = res_future.get();
399
  auto res2 = res_future2.get();
400
  // Assert the result
401
  ASSERT_VARIABLE_EQ(res, res_ref);
402
  ASSERT_VARIABLE_EQ(res2, res_ref);
403
}
404

405
TEST(TestRuntimeTaskAPI, TestTaskAPISameTasksMultiTensorInputs) {
406
  if (!torch_ipex::runtime::is_runtime_ext_enabled()) {
407
    GTEST_SKIP()
408
        << "Skip TestRuntimeTaskAPI::TestTaskAPISameTasksMultiTensorInputs. Didn't preload IOMP.";
409
  }
410
  std::vector<int32_t> cpu_core_list({0});
411
  torch_ipex::runtime::CPUPool cpu_pool(cpu_core_list);
412
  std::shared_ptr<torch_ipex::runtime::TaskExecutor> task_executor =
413
      std::make_shared<torch_ipex::runtime::TaskExecutor>(cpu_pool);
414

415
  at::Tensor input_tensor = at::rand({100, 8276});
416
  at::Tensor input_tensor2 = at::rand({100, 8276});
417
  // Get the reference result
418
  auto res_ref = at::softmax(input_tensor, -1);
419
  auto res_ref2 = at::softmax(input_tensor2, -1);
420
  // Create the task
421
  torch_ipex::runtime::
422
      Task<at::Tensor (*)(const at::Tensor&), const at::Tensor&>
423
          task(taskfunction_const_lvalue_reference, task_executor);
424

425
  auto res_future = task(input_tensor);
426
  auto res_future2 = task(input_tensor2);
427
  auto res = res_future.get();
428
  auto res2 = res_future2.get();
429
  // Assert the result
430
  ASSERT_VARIABLE_EQ(res, res_ref);
431
  ASSERT_VARIABLE_EQ(res2, res_ref2);
432
}
433