Update on "[inductor] Add AOTI ABI shim function for repeat_interleave.Tensor"

aakhundov · aakhundov · commit c6ffb8881bf3 · 2023-10-06T15:46:07.000-07:00
Summary: `repeat_interleave.Tensor` doesn't have inductor lowering. To invoke the operator in AOT Inductor's ABI compatibility mode we need a dedicated shim function.

Test Plan:

```
$ python test/inductor/test_aot_inductor.py -k test_repeat_interleave
...
----------------------------------------------------------------------
Ran 4 tests in 70.526s

OK
```

Reviewers:

Subscribers:

Tasks:

Tags:

cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler

[ghstack-poisoned]
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -698,9 +698,7 @@ def __init__(self):
             def forward(self, x):
                 return torch.ops.aten.repeat_interleave.Tensor(x, output_size=12)
 
-        example_inputs = (
-            torch.ones((1,), dtype=torch.int32, device="cuda") * 12,
-        )
+        example_inputs = (torch.ones((1,), dtype=torch.int32, device="cuda") * 12,)
         self.check_model(Repro(), example_inputs)
 
     def test_dynamic_cat(self):
diff --git a/torch/csrc/inductor/aoti_torch/shim_common.cpp b/torch/csrc/inductor/aoti_torch/shim_common.cpp
@@ -311,7 +311,8 @@ AOTITorchError aoti_torch_repeat_interleave_Tensor(
     AtenTensorHandle* out) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     at::Tensor* repeats_tensor = tensor_handle_to_tensor_pointer(repeats);
-    at::Tensor out_tensor = at::_ops::repeat_interleave_Tensor::call(*repeats_tensor, output_size);
+    at::Tensor out_tensor =
+        at::_ops::repeat_interleave_Tensor::call(*repeats_tensor, output_size);
     at::Tensor* out_tensor_ptr = new at::Tensor(std::move(out_tensor));
     *out = tensor_pointer_to_tensor_handle(out_tensor_ptr);
   });
