[super ugly maybe working code] use shim.h instead of Tensor

[janeyx99]

this is a poc change to see what making custom ops use aoti shim.h would look like

what you should expect by the end of this exercise:

• a libtorch.h that will contain ABI stable content only. the goal for this file is that it will go live in libtorch as a header only stable portion
• a libtorch.cpp that can call unstable things, because this will ultimately be upstreamed to core as a part of libtorch.so
• tensor_core_tiled_layout.cu will only depend on libtorch.h

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1548

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCM Infra failures during checkout of PyTorch

❌ 9 New Failures

As of commit 5e2c2d0 with merge base b2fb664 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 322979872e794b83a11f0e0a0e80a9a96d4644adc684bf9a297181f7b8a418cd /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 2a0a71f94bae3c24f0a8b855803b6a175c67c7f05cb8a80bd57c577e441f96ce /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 1368e1167e52bbfcba1fdfdc8869d4a6b5d180da4ca0ab868431d39fda5a2dbb /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t e5d71d9d541782120813ca1c18ad97bba95cf54f99b233bdbe43a744ded87496 /exec failed with exit code 139
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 117b1375e4ca6b83eaf92a735ce4bcfef158428a2f84708dd8f3d54c9bc0a2dc /exec failed with exit code 139
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  test/sparsity/test_marlin.py::SparseMarlin24::test_quant_sparse_marlin_layout_eager
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch==2.7.0.dev20250122 --index-url https://down... / linux-job (gh)
  RuntimeError: Command docker exec -t 23c2da5575a5aa9689d6631d84eeaedf87c8cdcd41a8afbd32d145d13b55a883 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch==2.7.0.dev20250122 --index-... / linux-job (gh)
  RuntimeError: Command docker exec -t b578215d4fd8e68ecec5cf1bdb48e38f740af6c2104ce98b70660a00d5f8eeaa /exec failed with exit code 139
• Run TorchAO Experimental Tests / test (macos-14) (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[kimishpatel]

are these, torch::jit:* and IValue, considered stable?

[janeyx99]

no! I haven't gotten a chance to document what I'm trying to do, but everything in libtorch.cpp should eventually make its way back to libtorch (not stable!!)

everything in libtorch.h is supposed to be stable.

[kimishpatel]

I am guessing that, from design perspective, this is the interface layer with registration API that is shipped with different versions of libtorch? But the user code, for custom ops, never relies on at::Tensor?

[janeyx99]

Yea, this is to allow for registering libtorch-agnostic custom ops which cannot use at::Tensor or IValue, etc. in their schema

[kimishpatel]

voidyvoid.. lol

[kimishpatel]

on serious note though, will user have to write this function with void** stack? It would actually be good if user just write with stable API. LIke AtenTensorHandle

[janeyx99]

Yea, to be clear, there are two portions we expect the user to provide for now:

1. their custom op, which should directly use AtenTensorHandle (we can make that easier by wrapping AtenTensorHandle with a headeronly C++ API layer)
2. how their custom op should be registered within our dispatcher stack (this is the point of the voidyvoid function). There may be a way to automatically generate this for the user given the schema of their custom op from (1) but that is a next step. Currently, we are going to expect users to provide this.