fix: revive torch backend (#13280)

* fix: revive torch backend

* as_strided view vs copy

* Revert "as_strided view vs copy"

This reverts commit 82a61223f2.

* add extra tests (move inplace, add fusion tests)

* better fusion with inplace_op

* no optimizer hooks (break mnist training fusion)

* split off fusion tests in separate file, assert on resnet fusion

fix: remove comments

* cleanup, reduce diff

* reduce diff

* better fusion and identity checks

---------

Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
This commit is contained in:
Roelof van Dijk
2025-11-20 00:26:50 +01:00
committed by GitHub
parent 56b2540349
commit 0dc2ff431d
6 changed files with 1140 additions and 232 deletions

View File

@@ -86,65 +86,67 @@ jobs:
clang -O2 recognize.c -lm -o recognize
cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
# TODO: fix the torch backend and reenable
# torchbackend:
# name: Torch Backend Tests
# runs-on: ubuntu-latest
# timeout-minutes: 15
# steps:
# - name: Checkout Code
# uses: actions/checkout@v4
# - name: Setup Environment
# uses: ./.github/actions/setup-tinygrad
# with:
# key: torch-backend-pillow-torchvision-et-pt
# deps: testing_minimal
# pydeps: "pillow torchvision expecttest"
# llvm: 'true'
# - name: Install ninja
# run: |
# sudo apt update || true
# sudo apt install -y --no-install-recommends ninja-build
# - name: Lint with ruff
# run: |
# pip3 install --upgrade --force-reinstall ruff==0.11.0
# python3 -m ruff check extra/torch_backend/backend.py
# - name: Test one op
# run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
# - name: Test ResNet-18
# run: DEBUG=2 python3 extra/torch_backend/example.py
# - name: My (custom) tests
# run: python3 extra/torch_backend/test.py
# - name: Test one op in torch tests
# run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
# - name: Test Ops with TINY_BACKEND
# run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
# - name: Test in-place operations on views
# run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
# - name: Test multi-gpu
# run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
torchbackend:
name: Torch Backend Tests
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_minimal
pydeps: "pillow torchvision expecttest"
llvm: 'true'
- name: Install ninja
run: |
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Lint with ruff
run: |
pip3 install --upgrade --force-reinstall ruff==0.11.0
python3 -m ruff check extra/torch_backend/backend.py
- name: Test one op
run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
- name: Test ResNet-18
run: DEBUG=2 python3 extra/torch_backend/example.py
- name: My (custom) tests
run: python3 extra/torch_backend/test.py
- name: Test one op in torch tests
run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
- name: Test Ops with TINY_BACKEND
run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
- name: Test in-place operations on views
run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
- name: Test multi-gpu
run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
- name: Test kernel fusion
run: python3 extra/torch_backend/test_kernel_fusion.py
# torchbackendmore:
# name: Torch Backend Tests More
# runs-on: ubuntu-latest
# timeout-minutes: 15
# steps:
# - name: Checkout Code
# uses: actions/checkout@v4
# - name: Setup Environment
# uses: ./.github/actions/setup-tinygrad
# with:
# key: torch-backend-pillow-torchvision-et-pt
# deps: testing_minimal
# llvm: 'true'
# - name: Install ninja
# run: |
# sudo apt update || true
# sudo apt install -y --no-install-recommends ninja-build
# - name: Test beautiful_mnist in torch with TINY_BACKEND
# run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
# - name: Test some torch tests (expect failure)
# run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
torchbackendmore:
name: Torch Backend Tests More
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_minimal
llvm: 'true'
- name: Install ninja
run: |
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
- name: Test beautiful_mnist in torch with TINY_BACKEND
run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
- name: Test some torch tests (expect failure)
run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
bepython:
name: Python Backend