Update on "[inductor] Add AOTI ABI shim function for torch.nonzero"

aakhundov · aakhundov · commit 510cba59dd59 · 2023-10-06T15:53:44.000-07:00
Summary: `torch.nonzero` doesn't have inductor lowering (yet). 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_zero_grid_with_unbacked_symbols
...
----------------------------------------------------------------------
Ran 4 tests in 78.650s

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/torch/csrc/inductor/aoti_runtime/model.h b/torch/csrc/inductor/aoti_runtime/model.h
@@ -84,13 +84,15 @@ class RAIIAtenTensorHandle {
 
   int64_t stride(int64_t d) {
     int64_t stride;
-    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_stride(handle_.get(), d, &stride));
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_stride(handle_.get(), d, &stride));
     return stride;
   }
 
   int64_t storage_offset() {
     int64_t storage_offset;
-    AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_get_storage_offset(handle_.get(), &storage_offset));
+    AOTI_TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_storage_offset(handle_.get(), &storage_offset));
     return storage_offset;
   }
 
diff --git a/torch/csrc/inductor/aoti_torch/c/shim.h b/torch/csrc/inductor/aoti_torch/c/shim.h
@@ -109,27 +109,20 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_sizes(
     int64_t** ret_sizes // returns borrowed reference
 );
 
-AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_size(
-    AtenTensorHandle tensor,
-    int64_t d,
-    int64_t* ret_size
-);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_size(AtenTensorHandle tensor, int64_t d, int64_t* ret_size);
 
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_strides(
     AtenTensorHandle tensor,
     int64_t** ret_strides // returns borrowed reference
 );
 
-AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_stride(
-    AtenTensorHandle tensor,
-    int64_t d,
-    int64_t* ret_stride
-);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_get_stride(AtenTensorHandle tensor, int64_t d, int64_t* ret_stride);
 
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_get_storage_offset(
     AtenTensorHandle tensor,
-    int64_t* ret_storage_offset
-);
+    int64_t* ret_storage_offset);
 
 // This function will create a new tensor object and its pointer is returned
 // through *out. The caller is responsible for wrapping the tensor pointer
@@ -215,9 +208,8 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mm_out(
     AtenTensorHandle self,
     AtenTensorHandle mat2);
 
-AOTI_TORCH_EXPORT AOTITorchError aoti_torch_nonzero(
-    AtenTensorHandle self,
-    AtenTensorHandle *out);
+AOTI_TORCH_EXPORT AOTITorchError
+aoti_torch_nonzero(AtenTensorHandle self, AtenTensorHandle* out);
 
 AOTI_TORCH_EXPORT AOTITorchError aoti_torch_repeat_interleave_Tensor(
     AtenTensorHandle repeats,
diff --git a/torch/csrc/inductor/aoti_torch/shim_common.cpp b/torch/csrc/inductor/aoti_torch/shim_common.cpp
@@ -118,8 +118,7 @@ AOTITorchError aoti_torch_get_sizes(
 AOTITorchError aoti_torch_get_size(
     AtenTensorHandle tensor,
     int64_t d,
-    int64_t* ret_size
-) {
+    int64_t* ret_size) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     at::Tensor* t = tensor_handle_to_tensor_pointer(tensor);
     *ret_size = t->size(d);
@@ -138,8 +137,7 @@ AOTITorchError aoti_torch_get_strides(
 AOTITorchError aoti_torch_get_stride(
     AtenTensorHandle tensor,
     int64_t d,
-    int64_t* ret_stride
-) {
+    int64_t* ret_stride) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     at::Tensor* t = tensor_handle_to_tensor_pointer(tensor);
     *ret_stride = t->stride(d);
@@ -148,8 +146,7 @@ AOTITorchError aoti_torch_get_stride(
 
 AOTITorchError aoti_torch_get_storage_offset(
     AtenTensorHandle tensor,
-    int64_t* ret_storage_offset
-) {
+    int64_t* ret_storage_offset) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     at::Tensor* t = tensor_handle_to_tensor_pointer(tensor);
     *ret_storage_offset = t->storage_offset();
@@ -340,7 +337,7 @@ AOTITorchError aoti_torch_mm_out(
 
 AOTITorchError aoti_torch_nonzero(
     AtenTensorHandle self,
-    AtenTensorHandle *out) {
+    AtenTensorHandle* out) {
   AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({
     at::Tensor* self_tensor = tensor_handle_to_tensor_pointer(self);
     at::Tensor out_tensor = at::nonzero(*self_tensor);
