[Intel GPU] xpu-ops codegen via backend whitelist

[ZhiweiYan-96]

Motivation

This PR intends to enhance the gen.py to allow generate codes for a specific backend.

XPU operators need be registered in an hand-written way currently. Developers have no chance to take the advantage of shared code to handle tensor meta setting (like strides, proxy output, structured kernels). Manually porting code is erro-prone and may lead to high maintaining efforts.

We intend to enhance the bckend_whitelist argument in gen.py to generate independent backend specific files. Then backend developers could directly use such generated code

Usage

The specialty of XPU ops is that they lie in third_pary/torch-xpu-ops, so it codegen resides in a separate process.

We use the following commands to trigger XPU codegen

python -m torchgen.gen --source-path path/to/yaml/of/xpu --install-dir build/aten/src/ATen/xpu --per-operator-headers --static-dispatch-backend --backend-whitelist=XPU

The diff lies at backend-whitelist=XPU. The backend-whitelist key is an existent argument in torchgen. But we observe that it can be hardly used for generating source files like Registexxx.cpp. We enhance the backend-white list argument in gen.py, where we could use it to generate backend only code. Like build/aten/src/ATen/xpu/RegisterXPU.cpp, build/aten/src/ATen/xpu/ops/as_strided_native.h

The input of gen.py are code templates and operators yaml. We share the same templates in aten. A simplified yaml lies in third_party/torch-xpu-ops, which only includes the supported xpu operators. This yaml is a copy-and-modify of native_functions.yaml. No extra entry is added, the format is same as the one in aten

Result

All operators headers are generated in build/aten/src/ATen/xpu/ops independently, which would not affect operators declared/defined by CPU/CUDA or any other backend. XPU operators only include headers in this folder.

Verification

• We add a ut for verifying xpu codgen in test_codegen.py and it could be triggered through
  python -m tools.test.test_codegen -k BackendWhitelist

• In third-party/torch-xpu-ops, we migrate some operators to structured kernels style, where they are registered through REGISTER_XPU_DISPATCH or TORCH_IMPL_FUNC. And it is tested with test_xpu.py

Others

As for the longer term, we would unify code generation system with other backend also, like cuda, mps.

Stack from ghstack (oldest at bottom):

• [Intel GPU] Add XPU into device list of copy_impl #129452
• -> [Intel GPU] xpu-ops codegen via backend whitelist #127865
• [Intel GPU] Dispatch Stub support #127860

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @chauhang @d4l3k @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @mcarilli @ptrblck @leslie-fang-intel @EikanWang @voznesenskym @penguinwu @Guobing-Chen @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @LucasLLC @MeetVadakkanchery @mhorowitz

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

❌ 39 Cancelled Jobs, 13 Unrelated Failures

As of commit b31d00f with merge base 3a18577 ():

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• inductor / cuda12.1-py3.10-gcc9-sm86 (gh)
• inductor / cuda12.1-py3.10-gcc9-sm86 / build (gh)
• inductor / cuda12.1-py3.12-gcc9-sm86 / build (gh)
• inductor / cuda12.4-py3.10-gcc9-sm86 (gh)
• inductor / cuda12.4-py3.10-gcc9-sm86 / build (gh)
• inductor / linux-jammy-cpu-py3.12-gcc11-inductor-halide (gh)
• inductor / linux-jammy-cpu-py3.12-gcc11-inductor-halide / build (gh)
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor (gh)
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / build (gh)
• inductor / rocm6.1-py3.8-inductor (gh)
• inductor / rocm6.1-py3.8-inductor / build (gh)
• inductor-periodic / cuda12.1-py3.10-gcc9-sm80 / build (gh)
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks (gh)
• inductor-periodic / cuda12.1-py3.10-gcc9-sm86-periodic-dynamo-benchmarks / build (gh)
• linux-binary-libtorch-cxx11-abi / libtorch-cpu-shared-with-deps-cxx11-abi-build / build (gh)
• linux-binary-libtorch-cxx11-abi / libtorch-cpu-shared-with-deps-cxx11-abi-test (gh)
• linux-binary-libtorch-pre-cxx11 / libtorch-cpu-shared-with-deps-pre-cxx11-build / build (gh)
• linux-binary-libtorch-pre-cxx11 / libtorch-cpu-shared-with-deps-pre-cxx11-test (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda11_8-build / build (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda11_8-split-build / build (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda11_8-split-test (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda11_8-test (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_1-build / build (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_1-split-build / build (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_1-split-test (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_1-test (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_4-build / build (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_4-split-build / build (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_4-split-test (gh)
• linux-binary-manywheel / manywheel-py3_8-cuda12_4-test (gh)
• Mac MPS / macos-py3-arm64-mps (gh)
• trunk / linux-focal-cuda11.8-py3.10-gcc9-experimental-split-build-test (gh)
• trunk / linux-focal-cuda12.4-py3.10-gcc9-experimental-split-build-test (gh)
• trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 (gh)
• trunk / linux-focal-rocm6.1-py3.8 (gh)
• trunk / linux-focal-rocm6.1-py3.8 / build (gh)
• trunk / macos-py3-arm64-mps (gh)
• trunk / pytorch-linux-focal-py3-clang9-android-ndk-r21e-build / build (default, 1, 1, linux.2xlarge) (gh)
• xpu / linux-jammy-xpu-py3.8 (gh)

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Mac MPS / macos-py3-arm64 / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / before-test / llm-retrieval (gh) (matched llm-retrieval rule in flaky-rules.json)
  ##[error]The operation was canceled.
• trunk / libtorch-linux-focal-cuda12.1-py3.7-gcc9-debug / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / libtorch-linux-focal-cuda12.4-py3.7-gcc9-debug / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / linux-focal-cuda11.8-py3.10-gcc9-experimental-split-build / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / linux-focal-cuda12.1-py3.10-gcc9-no-ops / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / linux-focal-cuda12.4-py3.10-gcc9-experimental-split-build / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / linux-focal-cuda12.4-py3.10-gcc9-no-ops / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / macos-py3-arm64 / build (gh) (similar failure)
  ##[error]The operation was canceled.
• trunk / win-vs2019-cpu-py3 / build (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.
• trunk / win-vs2019-cuda11.8-py3 / build (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.
• xpu / linux-jammy-xpu-py3.8 / build (gh) (similar failure)
  ##[error]The operation was canceled.

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

[EikanWang]

Please add UT.

[EikanWang]

I suppose the test cases should be part of test_codegen.py but not a dedicated test file.

[ZhiweiYan-96]

Yes, I have moved the ut into test_codegen.py

[EikanWang]

@ZhiweiYan-96 , please fix UT failures and we need to enable it ASAP.

[EikanWang]

Why do we need to add this yaml to the source code? I think it should be the data file of test folder. Why is it not be a part of test/xpu?

[ZhiweiYan-96]

Indeed, have moved into test/xpu, thanks for the comment